refactor the protocolerror class and usage

Author: Junchao Wu

Diff for: tchannel/tornado/connection.py

@@ -42,6 +42,7 @@
 from ..messages.common import FlagsType
 from ..messages.error import ErrorMessage
 from ..messages.types import Types
+from .message_factory import build_protocol_error
 from .message_factory import MessageFactory
 
 try:
@@ -218,7 +219,6 @@ def _loop(self):
                 # keep continue message in the list
                 # pop all other type messages including error message
 
-                result = {"type": context.message.message_type}
                 if (context.message.message_type in self.CALL_RES_TYPES and
                         context.message.flags == FlagsType.fragment):
                     # still streaming, keep it for record
@@ -227,27 +227,22 @@ def _loop(self):
                     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,
-                        )
+                    protocol_error = build_protocol_error(
+                        context.message,
+                        context.message_id,
                     )
-                    result['value'] = protocol_error
-                    future.set_result(result)
+                    future.set_exception(protocol_error)
                     continue
 
                 if response and future.running():
-                    result['value'] = response
-                    future.set_result(result)
+                    future.set_result(response)
                 continue
 
             log.warn('Unconsumed message %s', context)
 
     # Basically, the only difference between send and write is that send
     # sets up a Future to get the response. That's ideal for peers making
     # calls. Peers responding to calls must use write.
-
     def send(self, message, message_id=None):
         """Send the given message up the wire.
         Use this for messages which have a response message.
@@ -614,16 +609,7 @@ def send_request(self, request):
 
     def adapt_result(self, f, request, response_future):
         if f.exception():
-            # TODO unexpected exception
-            response_future.set_exception(
-                f.exception()
-            )
-            return
-
-        result = f.result()
-
-        if result['type'] == Types.ERROR:
-            protocol_error = result['value']
+            protocol_error = f.exception()
             protocol_error.tracing = request.tracing
             response_future.set_exception(
                 TChannelException(protocol_error.description)
@@ -633,8 +619,8 @@ def adapt_result(self, f, request, response_future):
                 EventType.after_receive_protocol_error,
                 protocol_error,
             )
-        elif result['type'] == Types.CALL_RES:
-            response = result['value']
+        else:
+            response = f.result()
             response.tracing = request.tracing
             response_future.set_result(response)
             # event: after_receive_response

Diff for: tchannel/tornado/data.py

@@ -318,7 +318,7 @@ def close_argstreams(self, force=False):
                 stream.close()
 
 
-class ProtocolError(object):
+class ProtocolError(Exception):
     """Object to represent protocol error message"""
 
     def __init__(

Diff for: tchannel/tornado/message_factory.py

@@ -43,6 +43,33 @@
 from .stream import InMemStream
 
 
+def build_raw_error_message(protocol_error):
+    """build protocol level error message based on Error object"""
+    message = ErrorMessage(
+        code=protocol_error.code,
+        tracing=Tracing(
+            protocol_error.tracing.span_id,
+            protocol_error.tracing.parent_span_id,
+            protocol_error.tracing.trace_id,
+            protocol_error.tracing.traceflags),
+        description=protocol_error.description,
+    )
+
+    return message
+
+
+def build_protocol_error(message, message_id=None):
+    """build protocol level error message based on Error object"""
+
+    error = ProtocolError(
+        code=message.code,
+        description=message.description,
+        id=message_id,
+    )
+
+    return error
+
+
 class MessageFactory(object):
     """Provide the functionality to decompose and recompose
     streaming messages.
@@ -58,31 +85,6 @@ def __init__(self, remote_host=None, remote_host_port=None):
         self.in_checksum = {}
         self.out_checksum = {}
 
-    def build_raw_error_message(self, protocol_error):
-        """build protocol level error message based on Error object"""
-        message = ErrorMessage(
-            code=protocol_error.code,
-            tracing=protocol_error.Tracing(
-                protocol_error.tracing.span_id,
-                protocol_error.tracing.parent_span_id,
-                protocol_error.tracing.trace_id,
-                protocol_error.tracing.traceflags),
-            description=protocol_error.description,
-        )
-
-        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,
-            description=message.description,
-            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.