Merge pull request #38 from jhjaggars/connection-refactor

Connection refactor

Author: matburt

receptor/connection.py

@@ -37,32 +37,36 @@ class BaseProtocol(asyncio.Protocol):
     def __init__(self, receptor, loop):
         self.receptor = receptor
         self.loop = loop
-        self.connection = None
+        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):
@@ -73,8 +77,7 @@ def connection_made(self, transport):
         self.loop.create_task(self.wait_greeting())
 
     def connection_lost(self, exc):
-        if self.connection is not None:
-            self.receptor.remove_connection(self.connection)
+        self.receptor.remove_connection(self)
 
     def data_received(self, data):
         logger.debug(data)
@@ -100,9 +103,11 @@ async def wait_greeting(self):
 
     def handle_handshake(self, data):
         self.greeted = True
-        self.connection = self.receptor.add_connection(data["id"], data.get("meta", {}), self)
-        self.loop.create_task(self.watch_queue(data["id"], self.transport))
-        self.loop.create_task(self.connection.message_handler(self.incoming_buffer))
+        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):
         msg = json.dumps({
@@ -125,7 +130,7 @@ def handle_handshake(self, data):
         super().handle_handshake(data)
         logger.debug("Received handshake from client with id %s, responding...", data["id"])
         self.send_handshake()
-        self.connection.send_route_advertisement()
+        self.receptor.send_route_advertisement()
 
 
 async def create_peer(receptor, loop, host, port):
@@ -155,7 +160,7 @@ def connection_lost(self, exc):
     def handle_handshake(self, data):
         super().handle_handshake(data)
         logger.debug("Received handshake from server with id %s", data["id"])
-        self.connection.send_route_advertisement()
+        self.receptor.send_route_advertisement()
 
 
 class BasicControllerProtocol(asyncio.Protocol):

receptor/protocol.py

@@ -8,8 +8,10 @@
 
 from .router import MeshRouter
 from .work import WorkManager
-from .connection import Connection
+from .messages import envelope, directive
+from . import exceptions
 
+RECEPTOR_DIRECTIVE_NAMESPACE = 'receptor'
 logger = logging.getLogger(__name__)
 
 
@@ -26,6 +28,7 @@ def __init__(self, config, node_id=None, router_cls=None,
         if not os.path.exists(self.base_path):
             os.makedirs(os.path.join(self.config.default_data_dir, self.node_id))
         self.connection_manifest_path = os.path.join(self.base_path, "connection_manifest")
+        self.buffer_mgr = self.config.components_buffer_manager
         self.stop = False
 
     def _find_node_id(self):
@@ -42,7 +45,7 @@ async def watch_expire(self):
         while True:
             current_manifest = self.get_connection_manifest()
             for connection in current_manifest:
-                buffer = self.config.components_buffer_manager.get_buffer_for_node(connection["id"], self)
+                buffer = self.buffer_mgr.get_buffer_for_node(connection["id"], self)
                 for ident, message in buffer:
                     message_actual = json.loads(message)
                     if "expire_time" in message_actual and message_actual['expire_time'] < time.time():
@@ -84,20 +87,26 @@ def update_connection_manifest(self, connection):
             manifest.append(dict(id=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:
+            for data in buf.get():
+                if "cmd" in data and data["cmd"] == "ROUTE":
+                    self.handle_route_advertisement(data)
+                else:
+                    await self.handle_message(data)
+            await asyncio.sleep(.1)
 
-    def add_connection(self, id_, meta, protocol_obj):
-        buffer_mgr = self.config.components_buffer_manager
-        conn = Connection(id_, meta, protocol_obj, buffer_mgr, self)
-        self.update_connections(conn)
-        return conn
+    def add_connection(self, protocol_obj):
+        self.update_connections(protocol_obj)
 
     def remove_connection(self, protocol_obj):
         notify_connections = []
@@ -110,11 +119,99 @@ def remove_connection(self, protocol_obj):
                 self.router.debug_router()
                 self.update_connection_manifest(connection_node)
             notify_connections += self.connections[connection_node]
-        for active_connection in notify_connections:
-            active_connection.send_route_advertisement(self.router.get_edges())
+        self.send_route_advertisement(self.router.get_edges())
 
     async def shutdown_handler(self):
         while True:
             if self.stop:
                 return
             await asyncio.sleep(1)
+
+    def handle_route_advertisement(self, data):
+        self.router.add_edges(data["edges"])
+        self.send_route_advertisement(data["edges"], data["seen"])
+
+    def send_route_advertisement(self, edges=None, seen=[]):
+        edges = edges or self.router.get_edges()
+        seen = set(seen)
+        logger.debug("Emitting Route Advertisements, excluding {}".format(seen))
+        destinations = set(self.connections) - seen
+        seens = list(seen | destinations | {self.node_id})
+
+        # TODO: This should be a broadcast call to the connection manager
+        for target in destinations:
+            buf = self.buffer_mgr.get_buffer_for_node(target, self)
+            try:
+                buf.push(json.dumps({
+                    "cmd": "ROUTE",
+                    "id": self.node_id,
+                    "capabilities": self.work_manager.get_capabilities(),
+                    "groups": self.config.node_groups,
+                    "edges": edges,
+                    "seen": seens
+                }).encode("utf-8"))
+            except exceptions.ReceptorBufferError as e:
+                logger.exception("Receptor Buffer Write Error broadcasting routes and capabilities: {}".format(e))
+                # TODO: This might should be a hard shutdown event
+            except Exception as e:
+                logger.exception("Error trying to broadcast routes and capabilities: {}".format(e))
+
+    async def handle_directive(self, outer_env):
+        try:
+            namespace, _ = outer_env.inner_obj.directive.split(':', 1)
+            if namespace == RECEPTOR_DIRECTIVE_NAMESPACE:
+                await directive.control(self.router, outer_env.inner_obj)
+            else:
+                # other namespace/work directives
+                await self.work_manager.handle(outer_env.inner_obj)
+        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(
+                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,
+                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(
+                receptor=self,
+                recipient=outer_env.inner_obj.sender,
+                payload=str(e),
+                in_response_to=outer_env.inner_obj.message_id,
+                serial=outer_env.inner_obj.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
+        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)

receptor/receptor.py

@@ -52,6 +52,14 @@ def find_edge(self, left, right):
                 return edge
         return None
 
+    def add_edges(self, edges):
+        for edge in edges:
+            existing_edge = self.find_edge(edge[0], edge[1])
+            if existing_edge and existing_edge[2] > edge[2]:
+                self.update_node(edge[0], edge[1], edge[2])
+            else:
+                self.register_edge(*edge)
+
     def register_edge(self, left, right, cost):
         if left != self.node_id:
             self._nodes.add(left)