add streaming logic to python tchannel inbound server

Author: Junchao Wu

tchannel/exceptions.py

@@ -53,6 +53,9 @@ class InvalidChecksumException(TChannelException):
     """Represent invalid checksum type in the message"""
     pass
 
+class StreamingException(TChannelException):
+    """Represent Streaming Message Exception"""
+    pass
 
 class InvalidErrorCodeException(TChannelException):
     """Represent Invalid Error Code exception"""

tchannel/handler.py

@@ -102,9 +102,14 @@ class TChannelRequest(object):
 
     def __init__(self, context, conn):
         self.message = context.message
-        self.header = getattr(self.message, 'arg_2', None)
-        self.body = getattr(self.message, 'arg_3', None)
-        self.method = getattr(self.message, 'arg_1', None)
+
+        try:
+            self.method = self.message.args[0]
+            self.header = self.message.args[1]
+            self.body = self.message.args[2]
+        except:
+            pass
+
         self.connection = conn
         self.context = context
         self.id = context.message_id
@@ -116,21 +121,27 @@ class TChannelResponse(object):
     """TChannel Response Wrapper"""
 
     __slots__ = ('_connection', '_request',
-                 'resp_msg', 'id')
+                 'resp_msg', 'id', 'body',
+                 'headers')
 
     def __init__(self, request, conn):
         self._connection = conn
         self._request = request
-        self.resp_msg = CallResponseMessage()
+        self.body = ""
+        self.headers = ""
         self.id = request.id
+        self.resp_msg = None
 
     def write(self, chunk):
         # build response message
-        self.resp_msg.arg_3 += chunk
+        self.body += chunk
 
     def finish(self):
+        # TODO add status code into arg_1 area
+        self.resp_msg = CallResponseMessage(
+            args=["", self.headers, self.body]
+        )
         self._connection.finish(self)
-        self.resp_msg = CallResponseMessage()
 
     def update_resp_id(self):
         self.id += 1

tchannel/messages/__init__.py

@@ -9,6 +9,8 @@
 from .init_response import InitResponseMessage, init_res_rw
 from .ping_request import PingRequestMessage, ping_req_rw
 from .ping_response import PingResponseMessage, ping_res_rw
+from .call_request_continue import call_req_c_rw
+from .call_response_continue import call_res_c_rw
 
 RW = {
     Types.CALL_REQ: call_req_rw,
@@ -17,15 +19,18 @@
     Types.INIT_REQ: init_req_rw,
     Types.INIT_RES: init_res_rw,
     Types.PING_REQ: ping_req_rw,
-    Types.PING_RES: ping_res_rw
+    Types.PING_RES: ping_res_rw,
+    Types.CALL_REQ_CONTINUE: call_req_c_rw,
+    Types.CALL_RES_CONTINUE: call_res_c_rw,
 }
 
-
 __all__ = [
     "RW",
     "ChecksumType",
     "CallRequestMessage",
+    "CallRequestContinueMessage",
     "CallResponseMessage",
+    "CallResponseContinueMessage",
     "ErrorMessage",
     "InitRequestMessage",
     "InitResponseMessage",

tchannel/messages/base.py

@@ -25,3 +25,15 @@ def __str__(self):
             str(self.__class__.__name__),
             ", ".join(attrs)
         )
+
+    def encode(self):
+        """ Encode all the strings in the msg using
+        encode type from common.py
+        """
+        pass
+
+    def decode(self):
+        """ Decode all the strings in the msg using
+        decode type from common.py
+        """
+        pass

tchannel/messages/call_continue.py

@@ -0,0 +1,89 @@
+from __future__ import absolute_import
+
+from .types import Types
+from .. import rw
+from . import common, ChecksumType
+from .base import BaseMessage
+
+
+class CallContinueMessage(BaseMessage):
+    """Represent a continuation of a call request (across multiple frames)."""
+    max_args_num = 3
+
+    __slots__ = (
+        'flags',
+        'checksum',
+        'args',
+    )
+
+    def __init__(
+        self,
+        flags=0,
+        checksum=None,
+        args=None,
+    ):
+        self.flags = flags
+        if checksum is not None:
+            checksum = common.ChecksumType.standardize(checksum)
+        self.checksum = checksum or \
+            (common.ChecksumType.none, None)
+
+        self.args = args or []
+
+    def encode(self):
+        self.args = list(map(
+            lambda arg: self.encode_str(arg),
+            self.args))
+
+    def decode(self):
+        self.args = list(map(
+            lambda arg: self.decode_str(arg),
+            self.args))
+
+    def fragment(self, space_left, fragment_msg):
+        """Streaming Message got fragmented based on
+        payload size. All the data within space_left
+        will be kept. All the rest will be shifted to
+        next fragment message.
+
+        :param space_left: space left for current frame
+        :param fragment_msg: the type is either
+        CallRequestMessage or CallResponseMessage
+
+
+        :return: None if there is space left
+        or next fragment message
+        """
+        new_args = []
+        for i, arg in enumerate(self.args):
+            if space_left >= 2: # 2bytes for size
+                space_left -= 2
+
+                if arg is not None:
+                    arg_length = len(arg)
+                    if space_left < arg_length:
+                        fragment_msg.args.append(arg[space_left+1:])
+                        new_args.append(arg[:space_left])
+                        space_left = 0
+                    else:
+                        new_args.append(arg)
+                        space_left -= arg_length
+                        if space_left <= 2:
+                            # boundary for arg
+                            fragment_msg.args.append("")
+                else:
+                    new_args.append(arg)
+            else:
+                for l in range(i, len(self.args)):
+                    fragment_msg.args.append(self.args[l])
+                break
+
+        self.args = new_args
+        if space_left >= 0 and len(fragment_msg.args) == 0:
+            # don't need to fragment any more
+            self.flags = 0x00
+            return None
+        else:
+            self.flags = 0x01
+            return fragment_msg
+

tchannel/messages/call_request.py

@@ -1,28 +1,20 @@
 from __future__ import absolute_import
 
-from .base import BaseMessage
+from .call_request_continue import CallRequestContinueMessage
 from .types import Types
 from .. import rw
-from . import common
+from . import common, ChecksumType
 
 
-class CallRequestMessage(BaseMessage):
+class CallRequestMessage(CallRequestContinueMessage):
     """Initiate an RPC call."""
     message_type = Types.CALL_REQ
 
     __slots__ = (
-        'flags',
         'ttl',
         'tracing',
-
         'service',
         'headers',
-
-        'checksum',
-
-        'arg_1',
-        'arg_2',
-        'arg_3',
     )
 
     def __init__(
@@ -33,42 +25,65 @@ def __init__(
         service=None,
         headers=None,
         checksum=None,
-        arg_1=None,
-        arg_2=None,
-        arg_3=None,
+        args=None,
     ):
-        self.flags = flags
+        super(CallRequestMessage, self).__init__(flags, checksum, args)
         self.ttl = ttl
         self.tracing = tracing or common.Tracing(0, 0, 0, 0)
         self.service = service or ''
         self.headers = dict(headers) if headers else {}
-        if checksum is not None:
-            checksum = common.ChecksumType.standardize(checksum)
-        self.checksum = checksum or \
-            (common.ChecksumType.none, None)
 
-        self.arg_1 = arg_1 or ''
-        self.arg_2 = arg_2 or ''
-        self.arg_3 = arg_3 or ''
+    def encode(self):
+        super(CallRequestMessage, self).encode()
+        self.service = self.service.encode(common.ENCODE_TYPE)
+        self.headers = dict(map(
+            lambda (k, v): self.encode_pair(k, v),
+            self.headers.iteritems()))
+
+    def decode(self):
+        super(CallRequestMessage, self).decode()
+        self.service = self.service.decode(common.DECODE_TYPE)
+        self.headers = dict(map(
+            lambda (k, v): self.decode_pair(k, v),
+            self.headers.iteritems()))
+
+    def get_meta_size(self):
+        size = 0
+        size += 1       # flags: 1
+        size += 4       # ttl: 4
+        size += 25      # tracing: 24 | traceflags: 1
+        size += 1       # service~1
+        size += len(self.service)
+        size += 1       # nh: 1
+        for k, v in self.headers:
+            size += 1       # hk~1
+            size += len(k)
+            size += 1       # hv~1
+            size += len(v)
+
+        size += 1       # csumtype: 1
+        size += 1 if self.checksum[0] != ChecksumType.none else 0
+        return size
 
+    def get_size(self):
+        size = self.get_meta_size()
+
+        for arg in self.args:
+            size += 2
+            size += len(arg)
+        return size
 
 call_req_rw = rw.instance(
     CallRequestMessage,
     ("flags", rw.number(1)),    # flags:1
     ("ttl", rw.number(4)),      # ttl:4
-
     ("tracing", common.tracing_rw),     # tracing:24
                                         # traceflags: 1
-
     ("service", rw.len_prefixed_string(rw.number(1))),  # service~1
     ("headers", rw.headers(             # nh:1 (hk~1 hv~1){nh}
         rw.number(1),
         rw.len_prefixed_string(rw.number(1))
     )),
-
     ("checksum", common.checksum_rw),   # csumtype:1 (csum:4){0, 1}
-
-    ("arg_1", rw.len_prefixed_string(rw.number(2), is_binary=True)),  # arg1~2
-    ("arg_2", rw.len_prefixed_string(rw.number(2), is_binary=True)),  # arg2~2
-    ("arg_3", rw.len_prefixed_string(rw.number(2), is_binary=True)),  # arg3~2
+    ("args", rw.args(rw.number(2))),    # [arg1~2, arg2~2, arg3~2]
 )

tchannel/messages/call_request_continue.py

@@ -1,11 +1,31 @@
 from __future__ import absolute_import
 
-from .call_request import CallRequestMessage
 from .types import Types
+from .. import rw
+from . import common
+from .call_continue import CallContinueMessage
 
 
-class CallRequestContinueMessage(CallRequestMessage):
+class CallRequestContinueMessage(CallContinueMessage):
     """Represent a continuation of a call request (across multiple frames)."""
     message_type = Types.CALL_REQ_CONTINUE
 
-    # TODO
+    def __init__(
+        self,
+        flags=0,
+        checksum=None,
+        args=None,
+    ):
+        super(CallRequestContinueMessage, self).__init__(flags, checksum, args)
+
+    def fragment(self, space_left):
+        fragment_msg = CallRequestContinueMessage()
+        return super(CallRequestContinueMessage, self).fragment(space_left, fragment_msg)
+
+
+call_req_c_rw = rw.instance(
+    CallRequestContinueMessage,
+    ("flags", rw.number(1)),    # flags:1
+    ("checksum", common.checksum_rw),   # csumtype:1 (csum:4){0, 1}
+    ("args", rw.args(rw.number(2))),    # [arg1~2, arg2~2, arg3~2]
+)