add protocol error object

throw TChannelException when client receives protocol error message

Author: Junchao Wu

tchannel/event.py

@@ -32,19 +32,8 @@
 
 @unique
 class EventType(IntEnum):
-    """Types to represent system events
+    """Types to represent system events"""
 
-    Events:
-
-    send_request: before client sends request
-
-    send_response: after server sends response
-
-    receive_request: after server receive request
-
-    receive_response: after client receive response
-
-    """
     before_send_request = 0x00
     after_send_request = 0x01
 
@@ -57,6 +46,9 @@ class EventType(IntEnum):
     before_receive_response = 0x30
     after_receive_response = 0x31
 
+    after_receive_protocol_error = 0x40
+    after_send_protocol_error = 0x41
+
 
 class EventHook(object):
     """provide all event hook interfaces
@@ -103,6 +95,14 @@ def after_receive_response(self, response):
         """Called after a ``CALL_RESP`` message is read."""
         pass
 
+    def after_receive_protocol_error(self, error):
+        """Called after a ''error'' message is read."""
+        pass
+
+    def after_send_protocol_error(self, error):
+        """Called after a ''error'' message is sent."""
+        pass
+
 
 class EventEmitter(object):
     def __init__(self):

tchannel/tornado/connection.py

@@ -36,19 +36,20 @@
 from ..event import EventType
 from ..exceptions import ConnectionClosedException
 from ..exceptions import InvalidErrorCodeException
+from ..exceptions import TChannelException
 from ..io import BytesIO
 from ..messages.common import PROTOCOL_VERSION
 from ..messages.common import FlagsType
 from ..messages.error import ErrorMessage
 from ..messages.types import Types
+from .data import ProtocolError
 from .message_factory import MessageFactory
 
 try:
     import tornado.queues as queues  # included in 4.2
 except ImportError:
     import toro as queues
 
-
 log = logging.getLogger('tchannel')
 
 
@@ -222,6 +223,17 @@ def _loop(self):
                     future = self._outstanding.get(context.message_id)
                 else:
                     future = self._outstanding.pop(context.message_id)
+
+                    if context.message.message_type == Types.ERROR:
+                        protocol_error = (
+                            self.response_message_factory.build_protocol_error(
+                                context.message,
+                                context.message_id,
+                            ))
+
+                        future.set_result(protocol_error)
+                        continue
+
                 if response and future.running():
                     future.set_result(response)
                 continue
@@ -368,7 +380,7 @@ def _extract_handshake_headers(self, message):
             )
 
         (self.remote_host,
-            self.remote_host_port) = message.host_port.rsplit(':', 1)
+         self.remote_host_port) = message.host_port.rsplit(':', 1)
         self.remote_host_port = int(self.remote_host_port)
         self.remote_process_name = message.process_name
         self.requested_version = message.version
@@ -593,19 +605,38 @@ def send_request(self, request):
         # TODO: fire before_receive_response
 
         def adapt_tracing(f):
-            # fetch the request tracing for response
-            f.result().tracing = request.tracing
-            response_future.set_result(f.result())
-            # event: receive_response
-            if self.tchannel:
-                self.tchannel.event_emitter.fire(
-                    EventType.after_receive_response,
-                    f.result(),
+            if not f.exception():
+                # fetch the request tracing for response
+                f.result().tracing = request.tracing
+
+                if isinstance(f.result(), ProtocolError):
+                    protocol_error = f.result()
+                    response_future.set_exception(
+                        TChannelException(protocol_error.message)
+                    )
+                    # event: after_receive_protocol_error
+                    if self.tchannel:
+                        self.tchannel.event_emitter.fire(
+                            EventType.after_receive_protocol_error,
+                            protocol_error,
+                        )
+                else:
+                    response = f.result()
+                    response_future.set_result(response)
+                    # event: after_receive_response
+                    if self.tchannel:
+                        self.tchannel.event_emitter.fire(
+                            EventType.after_receive_response,
+                            response,
+                        )
+            else:
+                # TODO unexpected exception
+                response_future.set_exception(
+                    f.exception()
                 )
 
         tornado.ioloop.IOLoop.current().add_future(
             future,
             adapt_tracing,
         )
-
         return response_future

tchannel/tornado/data.py

@@ -316,3 +316,22 @@ def close_argstreams(self, force=False):
         for stream in self.argstreams:
             if stream.auto_close or force:
                 stream.close()
+
+
+class ProtocolError(object):
+    """Object to represent protocol error message
+
+    Note: Internal use only.
+    """
+
+    def __init__(
+        self,
+        code,
+        message,
+        id=None,
+        tracing=None,
+    ):
+        self.code = code
+        self.tracing = tracing
+        self.id = id
+        self.message = message

tchannel/tornado/message_factory.py

@@ -34,8 +34,10 @@
 from ..messages.common import Tracing
 from ..messages.common import generate_checksum
 from ..messages.common import verify_checksum
+from ..messages.error import ErrorMessage
 from ..zipkin.annotation import Endpoint
 from ..zipkin.trace import Trace
+from .data import ProtocolError
 from .data import Request
 from .data import Response
 from .stream import InMemStream
@@ -56,6 +58,30 @@ def __init__(self, remote_host=None, remote_host_port=None):
         self.in_checksum = {}
         self.out_checksum = {}
 
+    def build_raw_error_message(self, err):
+        """build protocol level error message based on Error object"""
+        message = ErrorMessage(
+            code=err.code,
+            tracing=err.Tracing(err.tracing.span_id,
+                                err.tracing.parent_span_id,
+                                err.tracing.trace_id,
+                                err.tracing.traceflags),
+            message=err.message,
+        )
+
+        return message
+
+    def build_protocol_error(self, message, message_id=None):
+        """build protocol level error message based on Error object"""
+
+        error = ProtocolError(
+            code=message.code,
+            message=message.message,
+            id=message_id,
+        )
+
+        return error
+
     def build_raw_request_message(self, request, args, is_completed=False):
         """build protocol level message based on request and args.
 

tchannel/zipkin/zipkin_trace.py

@@ -39,34 +39,50 @@ def __init__(self, tchannel=None, dst=None):
             # to dst. By default it writes to stdout
             self.tracer = DebugTracer(dst)
 
-    def before_send_request(self, context):
-        if not context.tracing.traceflags:
+    def before_send_request(self, request):
+        if not request.tracing.traceflags:
             return
 
         ann = annotation.client_send()
-        context.tracing.annotations.append(ann)
+        request.tracing.annotations.append(ann)
 
-    def before_receive_request(self, context):
-        if not context.tracing.traceflags:
+    def before_receive_request(self, request):
+        if not request.tracing.traceflags:
             return
 
         ann = annotation.server_recv()
-        context.tracing.annotations.append(ann)
+        request.tracing.annotations.append(ann)
 
-    def after_send_response(self, context):
-        if not context.tracing.traceflags:
+    def after_send_response(self, response):
+        if not response.tracing.traceflags:
             return
 
         # send out a pair of annotations{server_recv, server_send} to zipkin
         ann = annotation.server_send()
-        context.tracing.annotations.append(ann)
-        self.tracer.record([(context.tracing, context.tracing.annotations)])
+        response.tracing.annotations.append(ann)
+        self.tracer.record([(response.tracing, response.tracing.annotations)])
 
-    def after_receive_response(self, context):
-        if not context.tracing.traceflags:
+    def after_receive_response(self, response):
+        if not response.tracing.traceflags:
             return
 
         # send out a pair of annotations{client_recv, client_send} to zipkin
         ann = annotation.client_recv()
-        context.tracing.annotations.append(ann)
-        self.tracer.record([(context.tracing, context.tracing.annotations)])
+        response.tracing.annotations.append(ann)
+        self.tracer.record([(response.tracing, response.tracing.annotations)])
+
+    def after_receive_protocol_error(self, error):
+        if not error.tracing.traceflags:
+            return
+
+        ann = annotation.client_recv()
+        error.tracing.annotations.append(ann)
+        self.tracer.record([(error.tracing, error.tracing.annotations)])
+
+    def after_send_protocol_error(self, error):
+        if not error.tracing.traceflags:
+            return
+
+        ann = annotation.server_send()
+        error.tracing.annotations.append(ann)
+        self.tracer.record([(error.tracing, error.tracing.annotations)])

tests/integration/test_client_server.py

@@ -24,8 +24,7 @@
 
 from tchannel import tcurl
 from tchannel.exceptions import ConnectionClosedException
-from tchannel.messages import Types
-from tchannel.messages.error import ErrorCode
+from tchannel.exceptions import TChannelException
 from tchannel.tornado import TChannel
 from tchannel.tornado.connection import StreamConnection
 from tchannel.tornado.data import Response
@@ -126,8 +125,7 @@ def test_endpoint_not_found(tchannel_server, call_response):
 
     hostport = 'localhost:%d' % (tchannel_server.port)
 
-    response = yield tchannel.request(hostport).send(InMemStream(),
-                                                     InMemStream(),
-                                                     InMemStream())
-    assert response.message_type == Types.ERROR
-    assert response.code == ErrorCode.bad_request
+    with pytest.raises(TChannelException):
+        yield tchannel.request(hostport).send(InMemStream(),
+                                              InMemStream(),
+                                              InMemStream())