Merge pull request #59 from jhjaggars/protocol-to-stream

defining the interface for a connection

Author: matburt

receptor/__main__.py

@@ -21,7 +21,7 @@ def main(args=None):
             'disable_existing_loggers': False,
             'formatters': {
                 'verbose': {
-                    'format': '{levelname} {asctime} {module} {message}',
+                    'format': '{levelname} {asctime} {node_id} {module} {message}',
                     'style': '{',
                 }
             },
@@ -40,6 +40,13 @@ def main(args=None):
         }
     )
 
+    def _f(record):
+        record.node_id = config.default_node_id
+        return True
+
+    for h in logging.getLogger('receptor').handlers:
+        h.addFilter(_f)
+
     try:
         config.go()
     except Exception:

receptor/config.py

@@ -110,18 +110,10 @@ def __init__(self, args=None):
         self.add_config_option(
             section='node',
             key='listen',
-            default_value='receptor://0.0.0.0:8888',
-            value_type='str',
+            default_value=['receptor://0.0.0.0:8888'],
+            value_type='list',
             hint='Set/override IP address and port to listen on. If not set here or in a config file, the default is receptor://0.0.0.0:8888.',
         )
-        self.add_config_option(
-            section='node',
-            key='websocket_listen',
-            default_value='',
-            value_type='str',
-            hint='Set IP address and port to listen on for websocket clients in the form ws[s]://addr:port. '
-                 'If not set here or in a config file then it is disabled',
-        )
         self.add_config_option(
             section='node',
             key='peers',
@@ -182,17 +174,9 @@ def __init__(self, args=None):
         self.add_config_option(
             section='controller',
             key='listen',
-            default_value='receptor://0.0.0.0:8888',
-            value_type='str',
-            hint='Set IP address and port to listen on. If not set here or in a config file, the default is receptor://0.0.0.0/0:8888.',
-        )
-        self.add_config_option(
-            section='controller',
-            key='websocket_listen',
-            default_value='',
-            value_type='str',
-            hint='Set IP address and port to listen on for websocket clients in the form ws://addr:port. '
-                 'If not set here or in a config file then it is disabled',
+            default_value=['receptor://0.0.0.0:8888'],
+            value_type='list',
+            hint='Set IP address and port to listen on. If not set here or in a config file, the default is receptor://0.0.0.0/0:8888. This option can be passed multiple times.',
         )
         self.add_config_option(
             section='controller',

receptor/connection/__init__.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/base.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/manager.py

@@ -0,0 +1,51 @@
+import asyncio
+import logging
+from .base import Transport
+
+logger = logging.getLogger(__name__)
+
+
+class RawSocket(Transport):
+    def __init__(self, reader, writer, chunk_size=2 ** 8):
+        self.reader = reader
+        self.writer = writer
+        self._closed = False
+        self.chunk_size = chunk_size
+
+    async def __anext__(self):
+        bytes_ = await self.reader.read(self.chunk_size)
+        return bytes_
+
+    @property
+    def closed(self):
+        return self._closed
+
+    async def close(self):
+        self._closed = True
+        await self.writer.close()
+
+    async def send(self, bytes_):
+        self.writer.write(bytes_)
+        await self.writer.drain()
+
+
+async def connect(host, port, factory, loop=None):
+    if not loop:
+        loop = asyncio.get_event_loop()
+
+    worker = factory()
+    try:
+        r, w = await asyncio.open_connection(host, port, loop=loop)
+        t = RawSocket(r, w)
+        await worker.client(t)
+    except Exception:
+        logger.exception("sock.connect")
+    finally:
+        await asyncio.sleep(5)
+        logger.debug("sock.connect: reconnection")
+        loop.create_task(connect(host, port, factory, loop))
+
+
+async def serve(reader, writer, factory):
+    t = RawSocket(reader, writer)
+    await factory().server(t)

receptor/connection/sock.py

@@ -0,0 +1,60 @@
+import logging
+import functools
+import aiohttp
+import aiohttp.web
+import asyncio
+
+from .base import Transport
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocket(Transport):
+    def __init__(self, ws):
+        self.ws = ws
+
+    async def __anext__(self):
+        msg = await self.ws.__anext__()
+        return msg.data
+
+    async def close(self):
+        return await self.ws.close()
+
+    @property
+    def closed(self):
+        return self.ws.closed
+
+    async def send(self, bytes_):
+        await self.ws.send_bytes(bytes_)
+
+
+async def connect(uri, factory, loop=None):
+    if not loop:
+        loop = asyncio.get_event_loop()
+
+    worker = factory()
+    try:
+        async with aiohttp.ClientSession().ws_connect(uri) as ws:
+            t = WebSocket(ws)
+            await worker.client(t)
+    except Exception:
+        logger.exception("ws.connect")
+    finally:
+        await asyncio.sleep(5)
+        logger.debug("ws.connect: reconnecting")
+        loop.create_task(connect(uri, factory=factory, loop=loop))
+
+
+async def serve(request, factory):
+    ws = aiohttp.web.WebSocketResponse()
+    await ws.prepare(request)
+
+    t = WebSocket(ws)
+    await factory().server(t)
+
+
+def app(factory):
+    handler = functools.partial(serve, factory=factory)
+    app = aiohttp.web.Application()
+    app.add_routes([aiohttp.web.get("/", handler)])
+    return app