moving connection management from controller to connection package

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

Author: jhjaggars

receptor/connection/__init__.py

@@ -1,122 +0,0 @@
-import logging
-
-import asyncio
-from collections.abc import AsyncIterator
-from abc import abstractmethod, abstractproperty
-
-from ..messages.envelope import FramedBuffer
-
-logger = logging.getLogger(__name__)
-
-
-class Transport(AsyncIterator):
-
-    @abstractmethod
-    async def close(self):
-        pass
-
-    @abstractproperty
-    def closed(self):
-        pass
-
-    @abstractmethod
-    async def send(self, bytes_):
-        pass
-
-
-async def watch_queue(conn, buf):
-    while not conn.closed:
-        try:
-            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
-
-        try:
-            await conn.send(msg)
-        except Exception:
-            logger.exception("watch_queue: error received trying to write")
-            await buf.put(msg)
-            return await conn.close()
-
-
-class Worker:
-    def __init__(self, receptor, loop):
-        self.receptor = receptor
-        self.loop = loop
-        self.conn = None
-        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):
-        self.read_task = self.loop.create_task(self.receive())
-
-    async def receive(self):
-        try:
-            async for msg in self.conn:
-                await self.buf.put(msg)
-        except Exception:
-            logger.exception("receive")
-
-    def register(self):
-        self.receptor.update_connections(self.conn, id_=self.remote_id)
-
-    def unregister(self):
-        self.receptor.remove_connection(self.conn, id_=self.remote_id, loop=self.loop)
-        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):
-        logger.debug("sending HI")
-        msg = self.receptor._say_hi().serialize()
-        await self.conn.send(msg)
-
-    async def start_processing(self):
-        logger.debug("sending routes")
-        await self.receptor.send_route_advertisement()
-        logger.debug("starting normal loop")
-        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
-        )
-        self.write_task = self.loop.create_task(watch_queue(self.conn, out))
-        return await self.write_task
-
-    async def _wait_handshake(self):
-        logger.debug("waiting for HI")
-        response = await self.buf.get()  # TODO: deal with timeout
-        self.remote_id = response.header["id"]
-        self.register()
-
-    async def client(self, transport):
-        try:
-            self.conn = transport
-            self.start_receiving()
-            await self.hello()
-            await self._wait_handshake()
-            await self.start_processing()
-            logger.debug("normal exit")
-        finally:
-            self.unregister()
-
-    async def server(self, transport):
-        try:
-            self.conn = transport
-            self.start_receiving()
-            await self._wait_handshake()
-            await self.hello()
-            await self.start_processing()
-        finally:
-            self.unregister()

receptor/connection/base.py

@@ -0,0 +1,121 @@
+import logging
+
+import asyncio
+from collections.abc import AsyncIterator
+from abc import abstractmethod, abstractproperty
+
+from ..messages.envelope import FramedBuffer
+
+logger = logging.getLogger(__name__)
+
+
+class Transport(AsyncIterator):
+    @abstractmethod
+    async def close(self):
+        pass
+
+    @abstractproperty
+    def closed(self):
+        pass
+
+    @abstractmethod
+    async def send(self, bytes_):
+        pass
+
+
+async def watch_queue(conn, buf):
+    while not conn.closed:
+        try:
+            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
+
+        try:
+            await conn.send(msg)
+        except Exception:
+            logger.exception("watch_queue: error received trying to write")
+            await buf.put(msg)
+            return await conn.close()
+
+
+class Worker:
+    def __init__(self, receptor, loop):
+        self.receptor = receptor
+        self.loop = loop
+        self.conn = None
+        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):
+        self.read_task = self.loop.create_task(self.receive())
+
+    async def receive(self):
+        try:
+            async for msg in self.conn:
+                await self.buf.put(msg)
+        except Exception:
+            logger.exception("receive")
+
+    def register(self):
+        self.receptor.update_connections(self.conn, id_=self.remote_id)
+
+    def unregister(self):
+        self.receptor.remove_connection(self.conn, id_=self.remote_id, loop=self.loop)
+        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):
+        logger.debug("sending HI")
+        msg = self.receptor._say_hi().serialize()
+        await self.conn.send(msg)
+
+    async def start_processing(self):
+        logger.debug("sending routes")
+        await self.receptor.send_route_advertisement()
+        logger.debug("starting normal loop")
+        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
+        )
+        self.write_task = self.loop.create_task(watch_queue(self.conn, out))
+        return await self.write_task
+
+    async def _wait_handshake(self):
+        logger.debug("waiting for HI")
+        response = await self.buf.get()  # TODO: deal with timeout
+        self.remote_id = response.header["id"]
+        self.register()
+
+    async def client(self, transport):
+        try:
+            self.conn = transport
+            self.start_receiving()
+            await self.hello()
+            await self._wait_handshake()
+            await self.start_processing()
+            logger.debug("normal exit")
+        finally:
+            self.unregister()
+
+    async def server(self, transport):
+        try:
+            self.conn = transport
+            self.start_receiving()
+            await self._wait_handshake()
+            await self.hello()
+            await self.start_processing()
+        finally:
+            self.unregister()

receptor/connection/manager.py

@@ -0,0 +1,42 @@
+import functools
+import asyncio
+from urllib.parse import urlparse
+
+from . import sock, ws
+
+
+def parse_peer(peer):
+    if "://" not in peer:
+        peer = f"receptor://{peer}"
+    return urlparse(peer)
+
+
+class Manager:
+    def __init__(self, factory, ssl_context, loop=None):
+        self.factory = factory
+        self.ssl_context = ssl_context
+        self.loop = loop or asyncio.get_event_loop()
+
+    def get_listener(self, listen_url):
+        service = parse_peer(listen_url)
+        if service.scheme == "receptor":
+            return asyncio.start_server(
+                functools.partial(sock.serve, factory=self.factory),
+                host=service.hostname,
+                port=service.port,
+                ssl=self.ssl_context,
+            )
+        elif service.scheme in ("ws", "wss"):
+            return self.loop.create_server(
+                ws.app(self.factory).make_handler(),
+                service.hostname,
+                service.port,
+                ssl=self.ssl_context,
+            )
+
+    def get_peer(self, peer):
+        service = parse_peer(peer)
+        if service.scheme == "receptor":
+            return sock.connect(service.hostname, service.port, self.factory, self.loop)
+        elif service.scheme in ("ws", "wss"):
+            return ws.connect(peer, self.factory, self.loop)

receptor/connection/sock.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from . import Transport
+from .base import Transport
 
 logger = logging.getLogger(__name__)
 

receptor/connection/ws.py

@@ -4,7 +4,7 @@
 import aiohttp.web
 import asyncio
 
-from . import Transport
+from .base import Transport
 
 logger = logging.getLogger(__name__)
 

receptor/controller.py

@@ -1,13 +1,12 @@
 import asyncio
 import datetime
-import functools
 import logging
 import uuid
-from urllib.parse import urlparse
 
 from .receptor import Receptor
 from .messages import envelope
-from .connection import ws, sock, Worker
+from .connection.base import Worker
+from .connection.manager import Manager
 
 logger = logging.getLogger(__name__)
 
@@ -17,45 +16,27 @@ class Controller:
     def __init__(self, config, loop=asyncio.get_event_loop(), queue=None):
         self.receptor = Receptor(config)
         self.loop = loop
-        self.factory = lambda: Worker(self.receptor, loop)
+        self.connection_manager = Manager(
+            lambda: Worker(self.receptor, loop),
+            self.receptor.config.get_server_ssl_context(),
+            loop
+        )
         self.queue = queue
         if self.queue is None:
             self.queue = asyncio.Queue(loop=loop)
         self.receptor.response_queue = self.queue
 
-    def parse_peer(self, peer):
-        if "://" not in peer:
-            peer = f"receptor://{peer}"
-        return urlparse(peer)
-
     def enable_server(self, listen_url):
-        service = self.parse_peer(listen_url)
-        client_connected_cb = functools.partial(sock.serve, factory=self.factory)
-        listener = asyncio.start_server(
-            client_connected_cb,
-            host=service.hostname,
-            port=service.port,
-            ssl=self.receptor.config.get_server_ssl_context())
-        logger.info("Serving on {}:{}".format(service.hostname, service.port))
+        listener = self.connection_manager.get_listener(listen_url)
+        logger.info("Serving on %s", listen_url)
         self.loop.create_task(listener)
 
     def enable_websocket_server(self, listen_url):
-        service = urlparse(listen_url)
-        listener = self.loop.create_server(
-            ws.app(self.factory).make_handler(),
-            service.hostname, service.port,
-            ssl=self.receptor.config.get_server_ssl_context())
-        logger.info("Serving websockets on {}:{}".format(service.hostname, service.port))
-        self.loop.create_task(listener)
+        self.enable_server(listen_url)
 
     async def add_peer(self, peer):
-        parsed = self.parse_peer(peer)
-        if parsed.scheme == 'receptor':
-            logger.info("Connecting to receptor peer {}".format(peer))
-            await self.loop.create_task(sock.connect(parsed.hostname, parsed.port, self.factory, self.loop))
-        elif parsed.scheme in ('ws', 'wss'):
-            logger.info("Connecting to websocket peer {}".format(peer))
-            await self.loop.create_task(ws.connect(peer, self.factory, self.loop))
+        logger.info("Connecting to peer {}".format(peer))
+        await self.loop.create_task(self.connection_manager.get_peer(peer))
 
     async def recv(self):
         inner = await self.receptor.response_queue.get()