Implement a better ping

Add coroutine support to the event loop

Author: matburt

receptor/controller.py

@@ -45,11 +45,11 @@ def enable_websocket_server(self, listen_address, listen_port):
         logger.info("Server ws on {}:{}".format(listen_address, listen_port))
         self.loop.create_task(listener)
 
-    def add_peer(self, peer):
+    async 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)))
+        await 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()
@@ -58,7 +58,7 @@ async def recv(self):
     async def send(self, message):
         inner_env = envelope.Inner(
             receptor=self.receptor,
-            messageid=str(uuid.uuid4()),
+            message_id=str(uuid.uuid4()),
             sender=self.receptor.node_id,
             recipient=message.recipient,
             message_type="directive",
@@ -71,9 +71,11 @@ async def send(self, message):
     async def ping(self, destination):
         await self.receptor.router.ping_node(destination)
 
-    def run(self):
+    def run(self, coro=None):
         try:
-            self.loop.run_until_complete(self.receptor.shutdown_handler())
+            if coro is None:
+                coro = self.receptor.shutdown_handler
+            self.loop.run_until_complete(coro())
         except KeyboardInterrupt:
             pass
         finally:

receptor/entrypoints.py

@@ -1,15 +1,11 @@
-import datetime
-import json
 import logging
-import sys
 import time
 import asyncio
 
 from prometheus_client import start_http_server
 
 from .receptor import Receptor
 from .controller import Controller
-from . import exceptions
 from . import node
 
 logger = logging.getLogger(__name__)
@@ -26,70 +22,46 @@ def run_as_node(config):
 
 def run_as_controller(config):
     controller = Controller(config)
+    if config.controller_stats_enable:
+        logger.info(f'Starting stats on port {config.node_stats_port}')
+        start_http_server(config.controller_stats_port)
     controller.enable_server(config.controller_listen_address,
                              config.controller_listen_port)
     controller.run()
 
 
 def run_as_ping(config):
+    def ping_iter():
+        if config.ping_count:
+            for x in range(config.ping_count):
+                yield x
+        else:
+            while True:
+                yield 0
+
+    async def handle_ping():
+        read_task = controller.loop.create_task(read_responses())
+        await controller.add_peer(config.ping_peer)
+        start_wait = time.time()
+        while not controller.receptor.router.node_is_known(config.ping_recipient) and (time.time() - start_wait < 5):
+            await asyncio.sleep(0.1)
+        await do_ping()
+        await read_task
+
     async def read_responses():
-        while True:
+        for _ in ping_iter():
             payload = await controller.recv()
             print("{}".format(payload))
 
     async def do_ping():
-        await asyncio.sleep(5)
-        await controller.ping(config.ping_recipient)
+        for _ in ping_iter():
+            await controller.ping(config.ping_recipient)
+            await asyncio.sleep(config.ping_delay)
 
     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)
+    controller.run(handle_ping)
 
-#     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):
     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()