Merge pull request #64 from matburt/fix_socket_thrashing

Fix socket thrashing on disconnect

Author: matburt

receptor/buffers/file.py

@@ -93,6 +93,7 @@ async def expire(self):
             while self.q.qsize() > 0:
                 ident = await self.q.get()
                 data = await self._get_file(ident, handle_only=True, delete=False)
+                # TODO: This will never work, it's not pure json anymore
                 msg = json.load(data)
                 if "expire_time" in msg and msg['expire_time'] < time.time():
                     logger.info("Expiring message %s", ident)

receptor/connection/base.py

@@ -27,6 +27,8 @@ async def watch_queue(conn, buf):
     while not conn.closed:
         try:
             msg = await asyncio.wait_for(buf.get(), 5.0)
+            if not msg:
+                return await conn.close()
         except asyncio.TimeoutError:
             continue
         except Exception:
@@ -58,6 +60,8 @@ def start_receiving(self):
     async def receive(self):
         try:
             async for msg in self.conn:
+                if self.conn.closed:
+                    break
                 await self.buf.put(msg)
         except Exception:
             logger.exception("receive")

receptor/connection/sock.py

@@ -14,15 +14,17 @@ def __init__(self, reader, writer, chunk_size=2 ** 8):
 
     async def __anext__(self):
         bytes_ = await self.reader.read(self.chunk_size)
+        if not bytes_:
+            self.close()
         return bytes_
 
     @property
     def closed(self):
         return self._closed
 
-    async def close(self):
+    def close(self):
         self._closed = True
-        await self.writer.close()
+        self.writer.close()
 
     async def send(self, bytes_):
         self.writer.write(bytes_)

receptor/entrypoints.py

@@ -45,8 +45,6 @@ def run_as_controller(config):
         logger.info(f'Starting stats on port {config.node_stats_port}')
         start_http_server(config.controller_stats_port)
     controller.enable_server(config.controller_listen)
-    if config.controller_websocket_listen:
-        controller.enable_websocket_server(config.controller_websocket_listen)
     controller.loop.create_task(controller.receptor.watch_expire())
     controller.run()
 

receptor/receptor.py

@@ -91,6 +91,7 @@ async def message_handler(self, buf):
                 data = await buf.get()
             except Exception:
                 logger.exception("message_handler")
+                break
             else:
                 logger.debug("message_handler: %s", data)
                 if "cmd" in data.header and data.header["cmd"] == "ROUTE":