Fixing up some bugs with disconnection handling

Author: matburt

receptor/controller.py

@@ -35,6 +35,7 @@ def mainloop(receptor, socket_path, loop=asyncio.get_event_loop()):
     )
     logger.info(f'Opening control socket on {socket_path}')
     loop.create_task(control_listener)
+    loop.create_task(receptor.watch_expire())
     try:
         loop.run_forever()
     except KeyboardInterrupt:

receptor/protocol.py

@@ -35,6 +35,11 @@ def __init__(self, receptor, loop):
         self.loop = loop
 
     async def watch_queue(self, node, transport):
+        '''
+        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():
@@ -63,13 +68,17 @@ def connection_made(self, transport):
         self.loop.create_task(self.wait_greeting())
 
     def connection_lost(self, exc):
-        self.receptor.remove_connection(self)
+        self.receptor.remove_connection(self.connection)
 
     def data_received(self, data):
         logger.debug(data)
         self.incoming_buffer.add(data)
 
     async def wait_greeting(self):
+        '''
+        Initialized when the connection is established to handle the greeting
+        before transitioning to message processing.
+        '''
         while not self.greeted:
             logger.debug('Looking for handshake...')
             for data in self.incoming_buffer.get():
@@ -80,9 +89,10 @@ async def wait_greeting(self):
                 else:
                     logger.error("Handshake failed!")
                     # TODO: Trigger disconnection
+                    # otherwise we could get stuck here
             await asyncio.sleep(.1)
         logger.debug("handshake complete, starting normal handle loop")
-        self.loop.create_task(self.connection.message_handler(self.incoming_buffer)) # Duplicated?
+        self.loop.create_task(self.connection.message_handler(self.incoming_buffer)) # Duplicated (see handle_handshake)?
 
     def handle_handshake(self, data):
         self.greeted = True

receptor/receptor.py

@@ -40,14 +40,12 @@ def _find_node_id(self):
 
     async def watch_expire(self):
         while True:
-            logger.info("Checking expirations")
             current_manifest = self.get_connection_manifest()
             for connection in current_manifest:
                 buffer = self.config.components.buffer_manager.get_buffer_for_node(connection["id"], self)
                 for ident, message in buffer:
                     message_actual = json.loads(message)
-                    logger.info("Examining {}".format(message_data))
-                    if "expire_time" in message_data and message_data['expire_time'] < time.time():
+                    if "expire_time" in message_actual and message_actual['expire_time'] < time.time():
                         logger.info("Expiring message {}:{}".format(ident, connection["id"]))
                         expired_message = buffer.read_message(ident, remove=True)
                         # TODO: Do something with expired message
@@ -83,7 +81,7 @@ def update_connection_manifest(self, connection):
                 found = True
                 break
         if not found:
-            manifest.append(dict(id=connection.id_,
+            manifest.append(dict(id=connection,
                             last=time.time()))
         self.write_connection_manifest(manifest)
 
@@ -102,19 +100,19 @@ def add_connection(self, id_, protocol_obj):
         return conn
 
     def remove_connection(self, protocol_obj):
-        notify_protocols = []
-        #self.update_connection_manifest(conn)
+        notify_connections = []
         for connection_node in self.connections:
             if protocol_obj in self.connections[connection_node]:
                 logger.info("Removing connection {} for node {}".format(protocol_obj, connection_node))
                 self.update_connection_manifest(connection_node)
                 self.connections[connection_node].remove(protocol_obj)
                 self.router.update_node(self.node_id, connection_node, 100)
                 self.router.debug_router()
-            notify_protocols += self.connections[connection_node]
+                self.update_connection_manifest(connection_node)
+            notify_connections += self.connections[connection_node]
         # TODO: Broadcast update, set timer for full expiration
-        for active_protocol in notify_protocols:
-            active_protocol.send_route_advertisement(self.router.get_edges())
+        for active_connection in notify_connections:
+            active_connection.send_route_advertisement(self.router.get_edges())
 
     async def shutdown_handler(self):
         while True: