Merge pull request #23 from jimi-c/ping_count

Make ping count/delay configurable

Author: matburt

receptor/config.py

@@ -181,6 +181,20 @@ def __init__(self, args=None):
             value_type='path',
             hint='Path to control socket for controller commands.',
         )
+        self.add_config_option(
+            section='ping',
+            key='count',
+            default_value=0,
+            value_type='int',
+            hint='Number of pings to send. If unspecified here or in a config file pings will be continuously sent until interrupted.',
+        )
+        self.add_config_option(
+            section='ping',
+            key='delay',
+            default_value=0,
+            value_type='float',
+            hint='The delay (in seconds) to wait between pings. If unspecified here or in a config file pings will be sent as soon as the previous response is received.',
+        )
         self.add_config_option(
             section='ping',
             key='recipient',
@@ -360,6 +374,8 @@ def _enforce_value_type(self, value, value_type):
                 return value_type(value)
             elif value_type == 'int' and not isinstance(value, int):
                 return int(value)
+            elif value_type == 'float' and not isinstance(value, float):
+                return float(value)
             elif value_type == 'str' and not isinstance(value, str):
                 return '%s' % (value,)
             elif value_type == 'bool' and not isinstance(value, bool):

receptor/connection.py

@@ -68,12 +68,24 @@ async def handle_message(self, msg):
         if next_hop is None:
             await outer_env.deserialize_inner(self.receptor)
             if outer_env.inner_obj.message_type == 'directive':
-                namespace, _ = outer_env.inner_obj.directive.split(':', 1)
-                if namespace == RECEPTOR_DIRECTIVE_NAMESPACE:
-                    await directive.control(self.receptor.router, outer_env.inner_obj)
-                else:
-                    # other namespace/work directives
-                    await self.receptor.work_manager.handle(outer_env.inner_obj)
+                try:
+                    namespace, _ = outer_env.inner_obj.directive.split(':', 1)
+                    if namespace == RECEPTOR_DIRECTIVE_NAMESPACE:
+                        await directive.control(self.receptor.router, outer_env.inner_obj)
+                    else:
+                        # other namespace/work directives
+                        await self.receptor.work_manager.handle(outer_env.inner_obj)
+                except Exception as e:
+                    err_resp = outer_env.inner_obj.make_response(
+                        receptor=self.receptor,
+                        recipient=outer_env.inner_obj.sender,
+                        payload=str(e),
+                        in_response_to=outer_env.inner_obj.message_id,
+                        serial=outer_env.inner_obj.serial + 1,
+                        ttl=15,
+                        code=1,
+                    )
+                    await self.receptor.router.send(err_resp)
             elif outer_env.inner_obj.message_type == 'response':
                 in_response_to = outer_env.inner_obj.in_response_to
                 if in_response_to in self.receptor.router.response_registry:
@@ -87,3 +99,4 @@ async def handle_message(self, msg):
                     f'Unknown message type: {outer_env.inner_obj.message_type}')
         else:
             await self.receptor.router.forward(outer_env, next_hop)
+

receptor/controller.py

@@ -1,25 +1,32 @@
 import asyncio
 import logging
+import os
 import socket
 import sys
-import os
 
 from . import protocol
 
 logger = logging.getLogger(__name__)
 
 
-def send_directive(directive, recipient, payload, socket_path):
-    if payload == '-':
-        payload = sys.stdin.read()
+def connect_to_socket(socket_path):
     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
     sock.connect(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:
-        response = sock.recv(4096)
-        sys.stdout.buffer.write(response + b"\n")
-        sys.stdout.flush()
+        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()):

receptor/entrypoints.py

@@ -1,5 +1,8 @@
 import datetime
+import json
 import logging
+import sys
+import time
 
 from .receptor import Receptor
 from . import node
@@ -14,18 +17,46 @@ def run_as_controller(config):
     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}')
+    node.mainloop(receptor, config.node_ping_interval)
+
+
 def run_as_ping(config):
     logger.info(f'Sending ping to {config.ping_recipient}.')
-    now = datetime.datetime.utcnow()
-    controller.send_directive('receptor:ping', config.ping_recipient, now.isoformat(), config.ping_socket_path)
+    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}.')
-    controller.send_directive(config.send_directive, config.send_recipient, config.send_payload, config.send_socket_path)
-
-
-def run_as_node(config):
-    receptor = Receptor(config)
-    logger.info(f'Running as Receptor node with ID: {receptor.node_id}')
-    node.mainloop(receptor, config.node_ping_interval)
+    sock = controller.connect_to_socket(config.send_socket_path)
+    try:
+        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()

receptor/protocol.py

@@ -1,10 +1,13 @@
-import datetime
 import asyncio
+import datetime
+import functools
+import json
 import logging
 import time
-import json
 import uuid
+
 from collections import deque
+
 from .messages import envelope
 from .exceptions import ReceptorBufferError
 
@@ -169,10 +172,14 @@ def connection_lost(self, exc):
             self.receptor.controller_connections.remove(self)
 
     def emit_response(self, response):
-        self.transport.write(json.dumps(dict(timestamp=response.timestamp,
-                                             in_response_to=response.in_response_to,
-                                             payload=response.raw_payload,
-                                             code=response.code)).encode())
+        emit_task = self.loop.create_task(response.sign_and_serialize())
+        emit_task.add_done_callback(
+            functools.partial(self._do_emit_callback)
+        )
+
+    def _do_emit_callback(self, fut):
+        res = fut.result()
+        self.transport.write(res.encode() + DELIM)
 
     def data_received(self, data):
         recipient, directive, payload = data.rstrip(DELIM).decode('utf8').split('\n', 2)
@@ -190,5 +197,28 @@ def data_received(self, data):
             directive=directive,
         )
         # TODO: Persistent registry?
-        self.loop.create_task(self.receptor.router.send(inner_env,
-                                                        expected_response=True))
+        send_task = self.loop.create_task(
+            self.receptor.router.send(
+                inner_env,
+                expected_response=True
+            )
+        )
+        send_task.add_done_callback(
+            functools.partial(self._data_received_callback, inner_env)
+        )
+
+    def _data_received_callback(self, inner_env, fut):
+        try:
+            fut.result()
+        except Exception as e:
+            err_resp = envelope.InnerEnvelope.make_response(
+                receptor=self.receptor,
+                recipient=inner_env.sender,
+                payload=str(e),
+                in_response_to=inner_env.message_id,
+                ttl=inner_env.ttl,
+                serial=inner_env.serial,
+                code=1,
+            )
+            self.emit_response(err_resp)
+