removing connection class

Signed-off-by: Jesse Jaggars <jjaggars@redhat.com>

Author: jhjaggars

receptor/connection.py

@@ -37,31 +37,36 @@ class BaseProtocol(asyncio.Protocol):
     def __init__(self, receptor, loop):
         self.receptor = receptor
         self.loop = loop
+        self.id = None
+        self.meta = None
 
-    async def watch_queue(self, node, transport):
+    def __str__(self):
+        return f"<Connection {self.id} {self.transport}"
+
+    async def watch_queue(self):
         '''
         Watches the buffer for this connection for messages delivered from other
         parts of Receptor (forwarded messages for example) for messages to send
         over the connection.
         '''
         buffer_mgr = self.receptor.config.components_buffer_manager
-        buffer_obj = buffer_mgr.get_buffer_for_node(node, self.receptor)
-        while not transport.is_closing():
+        buffer_obj = buffer_mgr.get_buffer_for_node(self.id, self.receptor)
+        while not self.transport.is_closing():
             try:
                 msg = buffer_obj.pop()
-                transport.write(msg + DELIM)
+                self.transport.write(msg + DELIM)
             except IndexError:
                 await asyncio.sleep(0.1)
             except ReceptorBufferError as e:
                 logger.exception("Receptor Buffer Read Error: {}".format(e))
                 # TODO: We need to try to send this message along somewhere else
                 # and record the failure somewhere
-                transport.close()
+                self.transport.close()
                 return
             except Exception as e:
-                logger.exception("Error received trying to write to {}: {}".format(node, e))
+                logger.exception("Error received trying to write to {}: {}".format(self.id, e))
                 buffer_obj.push(msg)
-                transport.close()
+                self.transport.close()
                 return
 
     def connection_made(self, transport):
@@ -98,8 +103,10 @@ async def wait_greeting(self):
 
     def handle_handshake(self, data):
         self.greeted = True
-        self.receptor.add_connection(data["id"], data.get("meta", {}), self)
-        self.loop.create_task(self.watch_queue(data["id"], self.transport))
+        self.id = data["id"]
+        self.meta = data.get("meta", {})
+        self.receptor.add_connection(self)
+        self.loop.create_task(self.watch_queue())
         self.loop.create_task(self.receptor.message_handler(self.incoming_buffer))
 
     def send_handshake(self):

receptor/protocol.py

@@ -89,13 +89,13 @@ def update_connection_manifest(self, connection):
                             last=time.time()))
         self.write_connection_manifest(manifest)
 
-    def update_connections(self, connection):
-        self.router.register_edge(connection.id_, self.node_id, 1)
-        if connection.id_ in self.connections:
-            self.connections[connection.id_].append(connection)
+    def update_connections(self, protocol_obj):
+        self.router.register_edge(protocol_obj.id, self.node_id, 1)
+        if protocol_obj.id in self.connections:
+            self.connections[protocol_obj.id].append(protocol_obj)
         else:
-            self.connections[connection.id_] = [connection]
-        self.update_connection_manifest(connection.id_)
+            self.connections[protocol_obj.id] = [protocol_obj]
+        self.update_connection_manifest(protocol_obj.id)
 
     async def message_handler(self, buf):
         while True:
@@ -106,9 +106,8 @@ async def message_handler(self, buf):
                     await self.handle_message(data)
             await asyncio.sleep(.1)
 
-    def add_connection(self, id_, meta, protocol_obj):
-        conn = Connection(id_, meta, protocol_obj)
-        self.update_connections(conn)
+    def add_connection(self, protocol_obj):
+        self.update_connections(protocol_obj)
 
     def remove_connection(self, protocol_obj):
         notify_connections = []
@@ -159,18 +158,7 @@ def send_route_advertisement(self, edges=None, seen=[]):
             except Exception as e:
                 logger.exception("Error trying to broadcast routes and capabilities: {}".format(e))
 
-    async def handle_message(self, msg):
-        outer_env = envelope.OuterEnvelope(**msg)
-        next_hop = self.router.next_hop(outer_env.recipient)
-        if next_hop:
-            return await self.router.forward(outer_env, next_hop)
-
-        await outer_env.deserialize_inner(self)
-
-        if outer_env.inner_obj.message_type != 'directive':
-            raise exceptions.UnknownMessageType(
-                f'Unknown message type: {outer_env.inner_obj.message_type}')
-
+    async def handle_directive(self, outer_env):
         try:
             namespace, _ = outer_env.inner_obj.directive.split(':', 1)
             if namespace == RECEPTOR_DIRECTIVE_NAMESPACE:
@@ -202,11 +190,30 @@ async def handle_message(self, msg):
                 code=1,
             )
             await self.router.send(err_resp)
-    elif outer_env.inner_obj.message_type == 'response':
+
+    async def handle_response(self, outer_env):
         in_response_to = outer_env.inner_obj.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)
         else:
             logger.warning(f'Received response to {in_response_to} but no record of sent message.')
+
+    async def handle_message(self, msg):
+        handlers = dict(
+            directive=self.handle_directive,
+            response=self.handle_response,
+        )
+        outer_env = envelope.OuterEnvelope(**msg)
+        next_hop = self.router.next_hop(outer_env.recipient)
+        if next_hop:
+            return await self.router.forward(outer_env, next_hop)
+
+        await outer_env.deserialize_inner(self)
+
+        if outer_env.inner_obj.message_type not in handlers:
+            raise exceptions.UnknownMessageType(
+                f'Unknown message type: {outer_env.inner_obj.message_type}')
+
+        await handlers[outer_env.inner_obj.message_type](outer_env)