using framed messages

jhjaggars · jhjaggars · commit 41e715947257 · 2019-11-21T17:57:27.000-06:00
Signed-off-by: Jesse Jaggars &lt;jjaggars@redhat.com&gt;
diff --git a/receptor/messages/envelope.py b/receptor/messages/envelope.py
@@ -5,15 +5,51 @@
 import itertools
 import json
 import logging
+import struct
 import time
 import uuid
-from struct import pack, unpack
+from enum import IntEnum
 
 logger = logging.getLogger(__name__)
 
 MAX_INT64 = (2 ** 64 - 1)
 
 
+class FramedMessage:
+    """
+    A complete message constructed from one or more Frames.
+    """
+
+    __slots__ = ("msg_id", "header", "payload")
+
+    def __init__(self, msg_id=None, header=None, payload=None):
+        if msg_id is None:
+            msg_id = uuid.uuid4().int
+        self.msg_id = msg_id
+        self.header = header
+        self.payload = payload
+
+
+    def serialize(self):
+        h = json.dumps(self.header).encode("utf-8")
+        return b''.join([
+            Frame.wrap(h, type_=Frame.Types.HEADER, msg_id=self.msg_id).serialize(),
+            h,
+            Frame.wrap(self.payload, msg_id=self.msg_id).serialize(),
+            self.payload
+        ])
+
+
+class CommandMessage(FramedMessage):
+
+    def serialize(self):
+        h = json.dumps(self.header).encode("utf-8")
+        return b''.join([
+            Frame.wrap(h, type_=Frame.Types.COMMAND, msg_id=self.msg_id).serialize(),
+            h,
+        ])
+
+
 class FramedBuffer:
     """
     A buffer that accumulates frames and bytes to produce a header and a
@@ -33,11 +69,11 @@ async def put(self, data):
         if not self.to_read:
             return await self.handle_frame(data)
         await self.consume(data)
-    
+
     async def handle_frame(self, data):
         self.current_frame, rest = Frame.from_data(data)
 
-        if self.current_frame.type not in (Frame.HEADER, Frame.PAYLOAD):
+        if self.current_frame.type not in Frame.Types:
             raise Exception("Unknown Frame Type")
 
         self.to_read = self.current_frame.length
@@ -50,11 +86,18 @@ async def consume(self, data):
             await self.finish()
 
     async def finish(self):
-        if self.current_frame.type == Frame.HEADER:
-            self.header = Header.from_bytes(self.bb)
-        elif self.current_frame.type == Frame.PAYLOAD:
-            await self.q.put((self.header, self.bb))
+        if self.current_frame.type == Frame.Types.HEADER:
+            self.header = json.loads(self.bb)
+        elif self.current_frame.type == Frame.Types.PAYLOAD:
+            await self.q.put(FramedMessage(
+                self.current_frame.msg_id, header=self.header,
+                payload=self.bb))
             self.header = None
+        elif self.current_frame.type == Frame.Types.COMMAND:
+            await self.q.put(FramedMessage(
+                self.current_frame.msg_id, header=json.loads(self.bb)))
+        else:
+            raise Exception("Unknown Frame Type")
         self.to_read = 0
         self.bb = bytearray()
 
@@ -63,11 +106,17 @@ async def get(self):
 
 
 class Frame:
-    HEADER = 0
-    PAYLOAD = 1
+
+    class Types(IntEnum):
+        HEADER = 0
+        PAYLOAD = 1
+        COMMAND = 2
+
+    fmt = struct.Struct(">ccIIQQ")
 
     __slots__ = ('type', 'version', 'length', 'msg_id', 'id')
 
+
     def __init__(self, type_, version, length, msg_id, id_):
         self.type = type_
         self.version = version
@@ -76,17 +125,29 @@ def __init__(self, type_, version, length, msg_id, id_):
         self.id = id_
 
     def serialize(self):
-        return pack(">ccIIQQ", bytes([self.type]), bytes([self.version]), self.id, self.length, *split_uuid(self.msg_id))
+        return self.fmt.pack(
+            bytes([self.type]), bytes([self.version]),
+            self.id, self.length, *split_uuid(self.msg_id))
 
     @classmethod
     def deserialize(cls, buf):
-        t, v, i, length, hi, lo = unpack(">ccIIQQ", buf)
+        t, v, i, length, hi, lo = Frame.fmt.unpack(buf)
         msg_id = join_uuid(hi, lo)
-        return cls(ord(t), ord(v), length, msg_id, i)
+        return cls(Frame.Types(ord(t)), ord(v), length, msg_id, i)
 
     @classmethod
     def from_data(cls, data):
-        return cls.deserialize(data[:26]), data[26:]
+        return cls.deserialize(data[:Frame.fmt.size]), data[Frame.fmt.size:]
+
+    @classmethod
+    def wrap(cls, data, type_=Types.PAYLOAD, msg_id=None):
+        """
+        Returns a frame for the passed data.
+        """
+        if not msg_id:
+            msg_id = uuid.uuid4().int
+
+        return cls(type_, 1, len(data), msg_id, 1)
 
 
 def split_uuid(data):
@@ -97,24 +158,10 @@ def join_uuid(hi, lo):
     return (hi << 64) | lo
 
 
-class Header:
-    def __init__(self, sender, recipient, route_list):
-        self.sender = sender
-        self.recipient = recipient
-        self.route_list = route_list
-
-    def serialize(self):
-        return json.dumps({"sender": self.sender, "recipient": self.recipient, "route_list": self.route_list}).encode("utf-8")
-
-    @classmethod
-    def from_bytes(cls, data):
-        return cls(**json.loads(data))
-
-    def __repr__(self):
-        return f"Header({self.sender}, {self.recipient}, {self.route_list})"
-
-    def __eq__(self, other):
-        return (self.sender, self.recipient, self.route_list) == (other.sender, other.recipient, other.route_list)
+def glue(header, payload):
+    hf = Frame.wrap(header, type_=Frame.Types.HEADER)
+    pf = Frame.wrap(payload, msg_id=hf.msg_id)
+    return hf.serialize() + header + pf.serialize() + payload
 
 
 def gen_chunks(data, header, msg_id=None, chunksize=2 ** 8):
@@ -123,9 +170,9 @@ def gen_chunks(data, header, msg_id=None, chunksize=2 ** 8):
     seq = itertools.count()
     buf = bytearray(chunksize)
     bv = memoryview(buf)
-    header = header.serialize()
-    yield Frame(Frame.HEADER, 1, len(header), msg_id, next(seq)).serialize() + header
-    yield Frame(Frame.PAYLOAD, 1, len(data), msg_id, next(seq)).serialize()
+    header = json.dumps(header).encode("utf-8")
+    yield Frame(Frame.Types.HEADER, 1, len(header), msg_id, next(seq)).serialize() + header
+    yield Frame(Frame.Types.PAYLOAD, 1, len(data), msg_id, next(seq)).serialize()
     buffer = io.BytesIO(data)
     bytes_read = buffer.readinto(buf)
     while bytes_read:
@@ -146,13 +193,13 @@ def __init__(self, frame_id, sender, recipient, route_list, inner):
         self.inner_obj = None
 
     async def deserialize_inner(self, receptor):
-        self.inner_obj = await InnerEnvelope.deserialize(receptor, self.inner)
+        self.inner_obj = await Inner.deserialize(receptor, self.inner)
 
     @classmethod
     def from_raw(cls, raw):
         doc = json.loads(raw)
         return cls(**doc)
-    
+
     def serialize(self):
         return json.dumps(dict(
             frame_id=self.frame_id,
@@ -163,7 +210,7 @@ def serialize(self):
         ))
 
 
-class InnerEnvelope:
+class Inner:
     def __init__(self, receptor, message_id, sender, recipient, message_type, timestamp,
                  raw_payload, directive=None, in_response_to=None, ttl=None, serial=1,
                  code=0, expire_time_delta=300):
diff --git a/receptor/protocol.py b/receptor/protocol.py
@@ -68,7 +68,7 @@ def connection_made(self, transport):
         self.peername = transport.get_extra_info('peername')
         self.transport = transport
         connected_peers_gauge.inc()
-        self.incoming_buffer = DataBuffer(loop=self.loop)
+        self.incoming_buffer = envelope.FramedBuffer(loop=self.loop)
         self.loop.create_task(self.wait_greeting())
 
     def connection_lost(self, exc):
@@ -77,7 +77,7 @@ def connection_lost(self, exc):
 
     def data_received(self, data):
         logger.debug(data)
-        self.incoming_buffer.add(data)
+        self.loop.create_task(self.incoming_buffer.put(data))
 
     async def wait_greeting(self):
         '''
@@ -86,9 +86,9 @@ async def wait_greeting(self):
         '''
         logger.debug('Looking for handshake...')
         data = await self.incoming_buffer.get()
-        logger.debug(data)
-        if data["cmd"] == "HI":
-            self.handle_handshake(data)
+        logger.debug(data.header)
+        if data.header["cmd"] == "HI":
+            self.handle_handshake(data.header)
             logger.debug("handshake complete, starting normal handle loop")
         else:
             logger.error("Handshake failed!")
@@ -102,15 +102,15 @@ def handle_handshake(self, data):
         self.loop.create_task(self.receptor.message_handler(self.incoming_buffer))
 
     def send_handshake(self):
-        msg = json.dumps({
+        msg = envelope.CommandMessage(header={
             "cmd": "HI",
             "id": self.receptor.node_id,
             "expire_time": time.time() + 10,
             "meta": dict(capabilities=self.receptor.work_manager.get_capabilities(),
                          groups=self.receptor.config.node_groups,
                          work=self.receptor.work_manager.get_work())
-        }).encode("utf-8")
-        self.transport.write(msg + DELIM)
+        })
+        self.transport.write(msg.serialize())
 
 
 class BasicProtocol(BaseProtocol):
diff --git a/receptor/receptor.py b/receptor/receptor.py
@@ -94,11 +94,11 @@ def update_connections(self, protocol_obj):
 
     async def message_handler(self, buf):
         while True:
-            data = await buf.get()
+            header, data = await buf.get()
             if "cmd" in data and data["cmd"] == "ROUTE":
                 await self.handle_route_advertisement(data)
             else:
-                await self.handle_message(data)
+                await self.handle_message(header, data)
 
     def add_connection(self, protocol_obj):
         self.update_connections(protocol_obj)
@@ -137,74 +137,74 @@ async def send_route_advertisement(self, edges=None, seen=[]):
         for target in destinations:
             buf = self.buffer_mgr.get_buffer_for_node(target, self)
             try:
-                await buf.put(json.dumps({
+                msg = envelope.CommandMessage(header={
                     "cmd": "ROUTE",
                     "id": self.node_id,
                     "capabilities": self.work_manager.get_capabilities(),
                     "groups": self.config.node_groups,
                     "edges": edges,
                     "seen": seens
-                }).encode("utf-8"))
+                })
+                await buf.put(msg.serialize())
             except Exception as e:
                 logger.exception("Error trying to broadcast routes and capabilities: {}".format(e))
 
-    async def handle_directive(self, outer_env):
+    async def handle_directive(self, inner):
         try:
-            namespace, _ = outer_env.inner_obj.directive.split(':', 1)
+            namespace, _ = inner.directive.split(':', 1)
             if namespace == RECEPTOR_DIRECTIVE_NAMESPACE:
-                await directive.control(self.router, outer_env.inner_obj)
+                await directive.control(self.router, inner)
             else:
                 # other namespace/work directives
-                await self.work_manager.handle(outer_env.inner_obj)
+                await self.work_manager.handle(inner)
         except ValueError:
-            logger.error("error in handle_message: Invalid directive -> '%s'. Sending failure response back." % (outer_env.inner_obj.directive,))
-            err_resp = outer_env.inner_obj.make_response(
+            logger.error("error in handle_message: Invalid directive -> '%s'. Sending failure response back." % (inner.directive,))
+            err_resp = inner.make_response(
                 receptor=self,
-                recipient=outer_env.inner_obj.sender,
-                payload="An invalid directive ('%s') was specified." % (outer_env.inner_obj.directive,),
-                in_response_to=outer_env.inner_obj.message_id,
-                serial=outer_env.inner_obj.serial + 1,
+                recipient=inner.sender,
+                payload="An invalid directive ('%s') was specified." % (inner.directive,),
+                in_response_to=inner.message_id,
+                serial=inner.serial + 1,
                 ttl=15,
                 code=1,
             )
             await self.router.send(err_resp)
         except Exception as e:
             logger.error("error in handle_message: '%s'. Sending failure response back." % (str(e),))
-            err_resp = outer_env.inner_obj.make_response(
+            err_resp = inner.make_response(
                 receptor=self,
-                recipient=outer_env.inner_obj.sender,
+                recipient=inner.sender,
                 payload=str(e),
-                in_response_to=outer_env.inner_obj.message_id,
-                serial=outer_env.inner_obj.serial + 1,
+                in_response_to=inner.message_id,
+                serial=inner.serial + 1,
                 ttl=15,
                 code=1,
             )
             await self.router.send(err_resp)
 
-    async def handle_response(self, outer_env):
-        in_response_to = outer_env.inner_obj.in_response_to
+    async def handle_response(self, inner):
+        in_response_to = inner.in_response_to
         if in_response_to in self.router.response_registry:
             logger.info(f'Handling response to {in_response_to} with callback.')
             for connection in self.controller_connections:
-                connection.emit_response(outer_env.inner_obj)
+                connection.emit_response(inner)
         else:
             logger.warning(f'Received response to {in_response_to} but no record of sent message.')
 
-    async def handle_message(self, msg):
+    async def handle_message(self, header, msg):
         handlers = dict(
             directive=self.handle_directive,
             response=self.handle_response,
         )
         messages_received_counter.inc()
-        outer_env = envelope.OuterEnvelope(**msg)
-        next_hop = self.router.next_hop(outer_env.recipient)
+        next_hop = self.router.next_hop(header.recipient)
         if next_hop:
-            return await self.router.forward(outer_env, next_hop)
+            return await self.router.forward(header, msg, next_hop)
 
-        await outer_env.deserialize_inner(self)
+        inner = await envelope.Inner.deserialize(self, msg)
 
-        if outer_env.inner_obj.message_type not in handlers:
+        if inner.message_type not in handlers:
             raise exceptions.UnknownMessageType(
-                f'Unknown message type: {outer_env.inner_obj.message_type}')
+                f'Unknown message type: {inner.message_type}')
 
-        await handlers[outer_env.inner_obj.message_type](outer_env)
+        await handlers[inner.message_type](inner)
diff --git a/receptor/router.py b/receptor/router.py
@@ -177,14 +177,14 @@ async def send(self, inner_envelope, expected_response=False):
             #TODO: This probably needs to emit an error response
             raise UnrouteableError(f'No route found to {inner_envelope.recipient}')
         signed = await inner_envelope.sign_and_serialize()
-        outer_envelope = envelope.OuterEnvelope(
-            frame_id=str(uuid.uuid4()),
-            sender=self.node_id,
-            recipient=inner_envelope.recipient,
-            route_list=[self.node_id],
-            inner=signed
-        )
+
+        header = {
+            "sender": self.node_id,
+            "recipient": inner_envelope.recipient,
+            "route_list": [self.node_id]
+        }
+        msg = envelope.FramedMessage(uuid.uuid4().int, header, signed)
         logger.debug(f'Sending {inner_envelope.message_id} to {inner_envelope.recipient} via {next_node_id}')
         if expected_response and inner_envelope.message_type == 'directive':
             self.response_registry[inner_envelope.message_id] = dict(message_sent_time=inner_envelope.timestamp)
-        await self.forward(outer_envelope, next_node_id)
+        await self.forward(msg, next_node_id)
diff --git a/receptor/tests/test_framedbuffer.py b/receptor/tests/test_framedbuffer.py
