Staging module changes

Author: matburt

receptor/__init__.py

@@ -0,0 +1,2 @@
+from .controller import Controller  # noqa
+from .config import ReceptorConfig  # noqa

receptor/__main__.py

@@ -8,7 +8,7 @@
 
 
 def main(args=None):
-    
+
     try:
         config = ReceptorConfig(args)
     except Exception as e:

receptor/config.py

@@ -5,7 +5,7 @@
 import os
 import ssl
 
-from .entrypoints import run_as_node, run_as_controller, run_as_ping, run_as_send
+from .entrypoints import run_as_node, run_as_ping, run_as_send, run_as_controller
 from .exceptions import ReceptorRuntimeError, ReceptorConfigError
 
 logger = logging.getLogger(__name__)
@@ -18,7 +18,7 @@
     },
     'controller': {
         'hint': 'Run a Receptor controller',
-        'entrypoint': run_as_controller,
+        'entrypoint': run_as_controller, # TODO: New entrypoint
     },
     'ping': {
         'hint': 'Tell the local controller to ping a node',
@@ -200,12 +200,19 @@ def __init__(self, args=None):
             hint='Set/override controller node identifier. If unspecified here or in a config file, one will be automatically generated.',
         )
         # Ping options
+        # self.add_config_option(
+        #     section='ping',
+        #     key='socket_path',
+        #     default_value='/var/run/receptor_controller.sock',
+        #     value_type='path',
+        #     hint='Path to control socket for controller commands.',
+        # )
         self.add_config_option(
             section='ping',
-            key='socket_path',
-            default_value='/var/run/receptor_controller.sock',
-            value_type='path',
-            hint='Path to control socket for controller commands.',
+            key='peer',
+            default_value='',
+            value_type='str',
+            hint='The peer to relay the ping directive through'
         )
         self.add_config_option(
             section='controller',

receptor/controller.py

@@ -1,11 +1,12 @@
 import asyncio
+import datetime
 import logging
-import os
-import socket
-import sys
+import uuid
 
-from . import protocol
 from .ws import WSServer
+from .protocol import BasicProtocol, create_peer
+from .receptor import Receptor
+from .messages import envelope
 
 logger = logging.getLogger(__name__)
 
@@ -16,47 +17,64 @@ def connect_to_socket(socket_path):
     return sock
 
 
-def send_directive(directive, recipient, payload, sock):
-    if payload == '-':
-        payload = sys.stdin.read()
-    sock.sendall(f"{recipient}\n{directive}\n{payload}".encode('utf-8') + protocol.DELIM)
-    response = b''
-    while True:
-        received = sock.recv(1024)
-        done = protocol.DELIM in received
-        response += received.rstrip(protocol.DELIM)
-        if done:
-            break
-            
-    return response
-
-
-# FIXME: the socket path is in the config, it shouldn't need to be passed as an arg here
-def mainloop(receptor, socket_path, loop=asyncio.get_event_loop()):
-    config = receptor.config
-    listener = loop.create_server(
-        lambda: protocol.BasicProtocol(receptor, 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
-    )
-    logger.info(f'Opening control socket on {socket_path}')
-    loop.create_task(control_listener)
-    loop.create_task(receptor.watch_expire())
-    try:
-        loop.run_forever()
-    except KeyboardInterrupt:
-        pass
-    finally:
-        loop.stop()
-        os.remove(socket_path)
+# TODO: track stats
+class Controller:
+
+    def __init__(self, config, loop=asyncio.get_event_loop(), queue=None):
+        self.receptor = Receptor(config)
+        self.loop = loop
+        self.queue = queue
+        if self.queue is None:
+            self.queue = asyncio.Queue(loop=loop)
+        self.receptor.response_queue = self.queue
+
+    def enable_server(self, listen_address, listen_port):
+        listener = self.loop.create_server(
+            lambda: BasicProtocol(self.receptor, self.loop),
+            listen_address, listen_port,
+            ssl=self.receptor.config.get_server_ssl_context())
+        logger.info("Serving on {}:{}".format(listen_address, listen_port))
+        # TODO: Enable stats?
+        self.loop.create_task(listener)
+
+    def enable_websocket_server(self, listen_address, listen_port):
+        listener = self.loop.create_server(
+            WSServer(self.receptor, self.loop).app().make_handler(),
+            listen_address, listen_port,
+            ssl=self.receptor.config.get_server_ssl_context())
+        logger.info("Server ws on {}:{}".format(listen_address, listen_port))
+        self.loop.create_task(listener)
+
+    def add_peer(self, peer):
+        # NOTE: Not signing or serializing
+        logger.info("Connecting to peer {}".format(peer))
+        self.loop.create_task(create_peer(self.receptor, self.loop,
+                                          *peer.strip().split(":", 1)))
+
+    async def recv(self):
+        inner = await self.receptor.response_queue.get()
+        return inner.raw_payload
+
+    async def send(self, message):
+        inner_env = envelope.Inner(
+            receptor=self.receptor,
+            messageid=str(uuid.uuid4()),
+            sender=self.receptor.node_id,
+            recipient=message.recipient,
+            message_type="directive",
+            directive=message.directive,
+            timestamp=datetime.datetime.utcnow().isoformat(),
+            raw_payload=message.fd.read(),
+        )
+        await self.receptor.router.send(inner_env, expected_response=True)
+
+    async def ping(self, destination):
+        await self.receptor.router.ping_node(destination)
+
+    def run(self):
+        try:
+            self.loop.run_until_complete(self.receptor.shutdown_handler())
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self.loop.stop()

receptor/entrypoints.py

@@ -3,26 +3,18 @@
 import logging
 import sys
 import time
+import asyncio
 
 from prometheus_client import start_http_server
 
 from .receptor import Receptor
-from . import controller
+from .controller import Controller
 from . import exceptions
 from . import node
 
 logger = logging.getLogger(__name__)
 
 
-def run_as_controller(config):
-    receptor = Receptor(config)
-    logger.info(f'Starting up as node ID {receptor.node_id}')
-    if config.controller_stats_enable:
-        logger.info(f'Starting stats on port {config.controller_stats_port}')
-        start_http_server(config.controller_stats_port)
-    controller.mainloop(receptor, config.controller_socket_path)
-
-
 def run_as_node(config):
     receptor = Receptor(config)
     logger.info(f'Running as Receptor node with ID: {receptor.node_id}')
@@ -32,47 +24,72 @@ def run_as_node(config):
     node.mainloop(receptor, config.node_ping_interval)
 
 
-def run_as_ping(config):
-    logger.info(f'Sending ping to {config.ping_recipient}.')
-    sock = controller.connect_to_socket(config.ping_socket_path)
+def run_as_controller(config):
+    controller = Controller(config)
+    controller.enable_server(config.controller_listen_address,
+                             config.controller_listen_port)
+    controller.run()
+
 
-    try:
-        pings_sent = 0
+def run_as_ping(config):
+    async def read_responses():
         while True:
-            now = datetime.datetime.utcnow()
-            response = controller.send_directive('receptor:ping', config.ping_recipient, now.isoformat(), sock)
-            resp_json = json.loads(response)
-            if 'code' in resp_json and resp_json['code'] != 0:
-                sys.stdout.buffer.write(b"Failed to ping node: %b\n" % (resp_json['raw_payload'].encode('utf-8'),))
-            else:
-                sys.stdout.buffer.write(response + b"\n")
-            sys.stdout.flush()
-            pings_sent += 1
-            if config.ping_count != 0 and pings_sent >= config.ping_count:
-                break
-            elif config.ping_delay > 0.0:
-                time.sleep(config.ping_delay)
-    except KeyboardInterrupt:
-        pass
-    finally:
-        sock.close()
+            payload = await controller.recv()
+            print("{}".format(payload))
+
+    async def do_ping():
+        await asyncio.sleep(5)
+        await controller.ping(config.ping_recipient)
+
+    logger.info(f'Sending ping to {config.ping_recipient} via {config.ping_peer}.')
+    controller = Controller(config)
+    controller.loop.create_task(read_responses())
+    controller.add_peer(config.ping_peer)
+    controller.loop.create_task(do_ping())
+    controller.run()
+    
+# def run_as_ping(config):
+#     logger.info(f'Sending ping to {config.ping_recipient}.')
+#     sock = controller.connect_to_socket(config.ping_socket_path)
 
+#     try:
+#         pings_sent = 0
+#         while True:
+#             now = datetime.datetime.utcnow()
+#             response = controller.send_directive('receptor:ping', config.ping_recipient, now.isoformat(), sock)
+#             resp_json = json.loads(response)
+#             if 'code' in resp_json and resp_json['code'] != 0:
+#                 sys.stdout.buffer.write(b"Failed to ping node: %b\n" % (resp_json['raw_payload'].encode('utf-8'),))
+#             else:
+#                 sys.stdout.buffer.write(response + b"\n")
+#             sys.stdout.flush()
+#             pings_sent += 1
+#             if config.ping_count != 0 and pings_sent >= config.ping_count:
+#                 break
+#             elif config.ping_delay > 0.0:
+#                 time.sleep(config.ping_delay)
+#     except KeyboardInterrupt:
+#         pass
+#     finally:
+#         sock.close()
 
 def run_as_send(config):
-    logger.info(f'Sending a {config.send_directive} directive to {config.send_recipient}.')
-    sock = controller.connect_to_socket(config.send_socket_path)
-    try:
-        if config.send_directive in (None, ''):
-            raise exceptions.UnknownDirective("The directive cannot be left out when using send.")
-        else:
-            try:
-                left, right = config.send_directive.split(':', 1)
-            except ValueError:
-                raise exceptions.UnknownDirective("Invalid directive format (%s). Directives must be in the form `action:method`." % (config.send_directive,))
-        response = controller.send_directive(config.send_directive, config.send_recipient, config.send_payload, sock)
-        sys.stdout.buffer.write(response + b"\n")
-        sys.stdout.flush()
-    except KeyboardInterrupt:
-        pass
-    finally:
-        sock.close()
+    pass
+# def run_as_send(config):
+#     logger.info(f'Sending a {config.send_directive} directive to {config.send_recipient}.')
+#     sock = controller.connect_to_socket(config.send_socket_path)
+#     try:
+#         if config.send_directive in (None, ''):
+#             raise exceptions.UnknownDirective("The directive cannot be left out when using send.")
+#         else:
+#             try:
+#                 left, right = config.send_directive.split(':', 1)
+#             except ValueError:
+#                 raise exceptions.UnknownDirective("Invalid directive format (%s). Directives must be in the form `action:method`." % (config.send_directive,))
+#         response = controller.send_directive(config.send_directive, config.send_recipient, config.send_payload, sock)
+#         sys.stdout.buffer.write(response + b"\n")
+#         sys.stdout.flush()
+#     except KeyboardInterrupt:
+#         pass
+#     finally:
+#         sock.close()