adding websockets

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

Author: jhjaggars

receptor/__main__.py

@@ -42,8 +42,8 @@ def main(args=None):
 
     try:
         config.go()
-    except Exception as e:
-        logger.error("An error occured while running receptor:\n%s" % (str(e),))
+    except Exception:
+        logger.exception("main: an error occured while running receptor")
         sys.exit(1)
 
 

receptor/controller.py

@@ -5,6 +5,7 @@
 import sys
 
 from . import protocol
+from .ws import WSServer
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +39,13 @@ def mainloop(receptor, socket_path, loop=asyncio.get_event_loop()):
         config.controller_listen_address, config.controller_listen_port, ssl=config.get_server_ssl_context())
     logger.info("Serving on %s:%s", config.controller_listen_address, config.controller_listen_port)
     loop.create_task(listener)
+
+    ws_server = WSServer(receptor, loop)
+    ws_listener = loop.create_server(ws_server.app().make_handler(),
+        config.node_listen_address, config.node_listen_port + 1, ssl=config.get_server_ssl_context())
+    loop.create_task(ws_listener)
+    logger.info("Serving ws on %s:%s", config.node_listen_address, config.node_listen_port + 1)
+
     control_listener = loop.create_unix_server(
         lambda: protocol.BasicControllerProtocol(receptor, loop),
         path=socket_path

receptor/messages/envelope.py

@@ -27,6 +27,9 @@ def __init__(self, msg_id=None, header=None, payload=None):
         self.header = header
         self.payload = payload
 
+    def __repr__(self):
+        return f"FramedMessage(msg_id={self.msg_id}, header={self.header}, payload={self.payload})"
+
     def serialize(self):
         h = json.dumps(self.header).encode("utf-8")
         return b''.join([
@@ -99,6 +102,7 @@ async def consume(self, data):
             await self.handle_frame(rest)
 
     async def finish(self):
+        logger.debug("finish: %s", self.current_frame)
         if self.current_frame.type == Frame.Types.HEADER:
             self.header = json.loads(self.bb)
         elif self.current_frame.type == Frame.Types.PAYLOAD:

receptor/node.py

@@ -1,10 +1,16 @@
 import asyncio
 import logging
+from urllib.parse import urlparse
 
 from .protocol import BasicProtocol, create_peer
+from .ws import WSClient, WSServer
 
 logger = logging.getLogger(__name__)
 
+def parse_peer(peer):
+    if "://" not in peer:
+        peer = f"receptor://{peer}"
+    return urlparse(peer)
 
 # FIXME: ping_interval is in the config, it shouldn't need to be passed as an arg here
 def mainloop(receptor, ping_interval=None, loop=asyncio.get_event_loop(), skip_run=False):
@@ -15,8 +21,21 @@ def mainloop(receptor, ping_interval=None, loop=asyncio.get_event_loop(), skip_r
             config.node_listen_address, config.node_listen_port, ssl=config.get_server_ssl_context())
         loop.create_task(listener)
         logger.info("Serving on %s:%s", config.node_listen_address, config.node_listen_port)
+
+        ws_server = WSServer(receptor, loop)
+        ws_listener = loop.create_server(ws_server.app().make_handler(),
+            config.node_listen_address, config.node_listen_port + 1, ssl=config.get_server_ssl_context())
+        loop.create_task(ws_listener)
+        logger.info("Serving ws on %s:%s", config.node_listen_address, config.node_listen_port + 1)
+
     for peer in config.node_peers:
-        loop.create_task(create_peer(receptor, loop, *peer.strip().split(":", 1)))
+        parsed = parse_peer(peer)
+        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))
+
     if ping_interval > 0:
         ping_time = (((int(loop.time()) + 1) // ping_interval) + 1) * ping_interval
         loop.call_at(ping_time, loop.create_task, send_pings_and_reschedule(receptor, loop, ping_time, ping_interval))

receptor/protocol.py

@@ -84,15 +84,7 @@ def handle_handshake(self, data):
         self.loop.create_task(self.receptor.message_handler(self.incoming_buffer))
 
     def send_handshake(self):
-        msg = envelope.CommandMessage(header={
-            "cmd": "HI",
-            "id": self.receptor.node_id,
-            "expire_time": time.time() + 10,
-            "meta": dict(capabilities=self.receptor.work_manager.get_capabilities(),
-                         groups=self.receptor.config.node_groups,
-                         work=self.receptor.work_manager.get_work())
-        })
-        self.transport.write(msg.serialize())
+        self.transport.write(self.receptor._say_hi().serialize())
 
 
 class BasicProtocol(BaseProtocol):

receptor/receptor.py

@@ -84,13 +84,15 @@ def update_connection_manifest(self, connection):
                             last=time.time()))
         self.write_connection_manifest(manifest)
 
-    def update_connections(self, protocol_obj):
-        self.router.register_edge(protocol_obj.id, self.node_id, 1)
-        if protocol_obj.id in self.connections:
-            self.connections[protocol_obj.id].append(protocol_obj)
+    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[protocol_obj.id] = [protocol_obj]
-        self.update_connection_manifest(protocol_obj.id)
+            self.connections[id_] = [protocol_obj]
+        self.update_connection_manifest(id_)
 
     async def message_handler(self, buf):
         logger.debug("spawning message_handler")
@@ -128,6 +130,16 @@ async def shutdown_handler(self):
                 return
             await asyncio.sleep(1)
 
+    def _say_hi(self):
+        return envelope.CommandMessage(header={
+            "cmd": "HI",
+            "id": self.node_id,
+            "expire_time": time.time() + 10,
+            "meta": dict(capabilities=self.work_manager.get_capabilities(),
+                         groups=self.config.node_groups,
+                         work=self.work_manager.get_work())
+        })
+
     async def handle_route_advertisement(self, data):
         self.router.add_edges(data["edges"])
         await self.send_route_advertisement(data["edges"], data["seen"])

receptor/ws.py

@@ -1,95 +1,104 @@
-import json
 import logging
-import time
 
+import asyncio
 import aiohttp
+import aiohttp.web
 
-from .protocol import DataBuffer
+from .messages.envelope import FramedBuffer
 
 logger = logging.getLogger(__name__)
 
 
-async def watch_queue(sock, buf):
-    while sock.open:
+async def watch_queue(ws, buf):
+    while not ws.closed:
         try:
             msg = await buf.get()
         except Exception:
-            logger.exception("Error getting data from buffer")
-        
+            logger.exception("watch_queue: error getting data from buffer")
+            continue
+
         try:
-            sock.send(msg)
+            await ws.send_bytes(msg)
         except Exception:
-            logger.exception("Error received trying to write")
+            logger.exception("watch_queue: error received trying to write")
             await buf.put(msg)
-            return await sock.close()
+            return await ws.close()
+    logger.debug("watch_queue: ws is now closed")
 
 
-class WSClient:
+class WSBase:
     def __init__(self, receptor, loop):
         self.receptor = receptor
         self.loop = loop
+        self.buf = FramedBuffer(loop=self.loop)
+        self.remote_id = None
 
-    async def connect(self, uri):
-        async with aiohttp.ClientSession().ws_connect(uri) as sock:
-            # handshake
-            node_id = await self.handshake(sock)
-            incoming_buffer = DataBuffer()
-            self.loop.create_task(self.receive(sock, incoming_buffer)) # reader
-
-            buf = self.receptor.buffer_mgr.get_buffer_for_node(node_id, self.receptor)
-            self.loop.create_task(watch_queue(sock, buf)) # writer
-
-        self.loop.create_task(self.connect(uri))
+    async def receive(self, ws):
+        try:
+            async for msg in ws:
+                await self.buf.put(msg.data)
+        except Exception:
+            logger.exception("receive")
 
+    def register(self, ws):
+        self.receptor.update_connections(ws, id_=self.remote_id)
 
-    async def handshake(self, sock):
-        msg = json.dumps({
-            "cmd": "HI",
-            "id": self.receptor.node_id,
-            "expire_time": time.time() + 10,
-            "meta": {
-                "capabilities": self.receptor.work_manager.get_capabilities(),
-                "groups": self.receptor.config.node_groups,
-                "work": self.receptor.work_manager.get_work(),
-            }
-        }).encode("utf-8")
-        await sock.send_bytes(msg)
-        response = await sock.receive().json()
-        return response["id"]
+    async def hello(self, ws):
+        msg = self.receptor._say_hi().serialize()
+        await ws.send_bytes(msg)
 
-    async def receive(self, sock, buf):
-        self.loop.create_task(self.receptor.message_handler(buf))
-        async for msg in sock.receive():
-            buf.add(msg.data)
+    async def start_processing(self, ws):
+        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)
 
 
-class WSServer:
+class WSClient(WSBase):
+    async def connect(self, uri):
+        try:
+            async with aiohttp.ClientSession().ws_connect(uri) as ws:
+
+                logger.debug("connect: starting recv")
+                recv_loop = self.loop.create_task(self.receive(ws))  # reader
+                logger.debug("connect: sending HI")
+                await self.hello(ws)
+                logger.debug("connect: waiting for HI")
+                response = await self.buf.get()  # TODO: deal with timeout
+                self.remote_id = response.header["id"]
+                self.register(ws)
+                logger.debug("connect: sending routes")
+                await self.receptor.send_route_advertisement()
+                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))
 
-    def __init__(self, receptor, loop):
-        self.receptor = receptor
-        self.loop = loop
 
+class WSServer(WSBase):
     async def serve(self, request):
 
         ws = aiohttp.web.WebSocketResponse()
         await ws.prepare(request)
 
-        handshake = await ws.receive().json()
-        await ws.send_json({
-            "cmd": "HI",
-            "id": self.receptor.node_id,
-            "expire_time": time.time() + 10,
-            "meta": {
-                "capabilities": self.receptor.work_manager.get_capabilities(),
-                "groups": self.receptor.config.node_groups,
-                "work": self.receptor.work_manager.get_work(),
-            }
-        })
-
-        buf = self.receptor.buffer_mgr.get_buffer_for_node(handshake["id"], self.receptor)
-        self.loop.create_task(watch_queue(ws, buf)) # writer
-
-        incoming_buffer = DataBuffer()
-        self.loop.create_task(self.receptor.message_handler(incoming_buffer))
-        async for msg in ws:
-            incoming_buffer.add(msg.data)
+        logger.debug("serve: starting recv")
+        self.loop.create_task(self.receive(ws))  # reader
+        logger.debug("serve: waiting for HI")
+        response = await self.buf.get()  # TODO: deal with timeout
+        self.remote_id = response.header["id"]
+        self.register(ws)
+        logger.debug("serve: sending HI")
+        await self.hello(ws)
+        logger.debug("serve: sending routes")
+        await self.receptor.send_route_advertisement()
+        logger.debug("serve: starting normal recv loop")
+        await self.start_processing(ws)
+
+    def app(self):
+        app = aiohttp.web.Application()
+        app.add_routes([aiohttp.web.get("/", self.serve)])
+        return app