fixing protocol selection and adding task cancellation

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

Author: jhjaggars

receptor/node.py

@@ -30,11 +30,13 @@ def mainloop(receptor, ping_interval=None, loop=asyncio.get_event_loop(), skip_r
 
     for peer in config.node_peers:
         parsed = parse_peer(peer)
-        if parsed.scheme == "receptor://":
+        if parsed.scheme == "receptor":
             loop.create_task(create_peer(receptor, loop, parsed.hostname, parsed.port))
         elif parsed.scheme in ("ws", "wss"):
             c = WSClient(receptor, loop)
             loop.create_task(c.connect(peer))
+        else:
+            print(f"invalid peer: {peer} -> {parsed}")
 
     if ping_interval > 0:
         ping_time = (((int(loop.time()) + 1) // ping_interval) + 1) * ping_interval
@@ -53,6 +55,6 @@ async def send_pings_and_reschedule(receptor, loop, ping_time, ping_interval):
     logger.debug(f'Scheduling mesh ping.')
     for node_id in receptor.router.get_nodes():
         await receptor.router.ping_node(node_id)
-    loop.call_at(ping_time + ping_interval, 
+    loop.call_at(ping_time + ping_interval,
                  loop.create_task, send_pings_and_reschedule(
                      receptor, loop, ping_time + ping_interval, ping_interval))

receptor/receptor.py

@@ -84,16 +84,6 @@ def update_connection_manifest(self, connection):
                             last=time.time()))
         self.write_connection_manifest(manifest)
 
-    def update_connections(self, protocol_obj, id_=None):
-        if id_ is None:
-            id_ = protocol_obj.id
-        self.router.register_edge(id_, self.node_id, 1)
-        if id_ in self.connections:
-            self.connections[id_].append(protocol_obj)
-        else:
-            self.connections[id_] = [protocol_obj]
-        self.update_connection_manifest(id_)
-
     async def message_handler(self, buf):
         logger.debug("spawning message_handler")
         while True:
@@ -108,10 +98,24 @@ async def message_handler(self, buf):
                 else:
                     await self.handle_message(data)
 
+    def update_connections(self, protocol_obj, id_=None):
+        if id_ is None:
+            id_ = protocol_obj.id
+
+        self.router.register_edge(id_, self.node_id, 1)
+        if id_ in self.connections:
+            self.connections[id_].append(protocol_obj)
+        else:
+            self.connections[id_] = [protocol_obj]
+        self.update_connection_manifest(id_)
+
     def add_connection(self, protocol_obj):
         self.update_connections(protocol_obj)
 
-    def remove_connection(self, protocol_obj):
+    def remove_connection(self, protocol_obj, id_=None):
+        if id_ is None:
+            id_ = protocol_obj.id
+
         notify_connections = []
         for connection_node in self.connections:
             if protocol_obj in self.connections[connection_node]:

receptor/ws.py

@@ -12,7 +12,9 @@
 async def watch_queue(ws, buf):
     while not ws.closed:
         try:
-            msg = await buf.get()
+            msg = await asyncio.wait_for(buf.get(), 5.0)
+        except asyncio.TimeoutError:
+            continue
         except Exception:
             logger.exception("watch_queue: error getting data from buffer")
             continue
@@ -32,6 +34,12 @@ def __init__(self, receptor, loop):
         self.loop = loop
         self.buf = FramedBuffer(loop=self.loop)
         self.remote_id = None
+        self.read_task = None
+        self.handle_task = None
+        self.write_task = None
+
+    def start_receiving(self, ws):
+        self.read_task = self.loop.create_task(self.receive(ws))
 
     async def receive(self, ws):
         try:
@@ -43,25 +51,35 @@ async def receive(self, ws):
     def register(self, ws):
         self.receptor.update_connections(ws, id_=self.remote_id)
 
+    def unregister(self, ws):
+        self.receptor.remove_connection(ws, id_=self.remote_id)
+        self._cancel(self.read_task)
+        self._cancel(self.handle_task)
+        self._cancel(self.write_task)
+
+    def _cancel(self, task):
+        if task:
+            task.cancel()
+
     async def hello(self, ws):
         msg = self.receptor._say_hi().serialize()
         await ws.send_bytes(msg)
 
     async def start_processing(self, ws):
-        self.loop.create_task(self.receptor.message_handler(self.buf))
+        self.handle_task = self.loop.create_task(self.receptor.message_handler(self.buf))
         out = self.receptor.buffer_mgr.get_buffer_for_node(
             self.remote_id, self.receptor
         )
-        return await watch_queue(ws, out)
+        self.write_task = self.loop.create_task(watch_queue(ws, out))
+        return await self.write_task
 
 
 class WSClient(WSBase):
     async def connect(self, uri):
-        try:
-            async with aiohttp.ClientSession().ws_connect(uri) as ws:
-
+        async with aiohttp.ClientSession().ws_connect(uri) as ws:
+            try:
                 logger.debug("connect: starting recv")
-                recv_loop = self.loop.create_task(self.receive(ws))  # reader
+                self.start_receiving(ws)
                 logger.debug("connect: sending HI")
                 await self.hello(ws)
                 logger.debug("connect: waiting for HI")
@@ -73,10 +91,13 @@ async def connect(self, uri):
                 logger.debug("connect: starting normal loop")
                 await self.start_processing(ws)
                 logger.debug("connect: normal exit")
-        except Exception:
-            logger.exception("connect")
-            logger.debug("connect: reconnecting")
-            self.loop.create_task(self.connect(uri))
+            except Exception:
+                logger.exception("connect")
+            finally:
+                self.unregister(ws)
+                await asyncio.sleep(5)
+                logger.debug("connect: reconnecting")
+                self.loop.create_task(self.connect(uri))
 
 
 class WSServer(WSBase):
@@ -97,6 +118,7 @@ async def serve(self, request):
         await self.receptor.send_route_advertisement()
         logger.debug("serve: starting normal recv loop")
         await self.start_processing(ws)
+        self.unregister(ws)
 
     def app(self):
         app = aiohttp.web.Application()