Merge pull request #35 from matburt/grouping_and_metadata

Grouping and metadata

Author: matburt

Makefile

@@ -17,7 +17,7 @@ else
     endif
 endif
 
-.PHONY: dist sdist
+.PHONY: dist sdist test
 
 clean:
 	rm -rf dist

receptor/config.py

@@ -144,6 +144,16 @@ def __init__(self, args=None):
             value_type='int',
             hint='If specified, the node will ping all other known nodes in the mesh every N seconds. The default is -1, meaning no pings are sent.',
         )
+        self.add_config_option(
+            section='node',
+            key='groups',
+            short_option='-g',
+            long_option='--group',
+            default_value=[],
+            value_type='list',
+            listof='str',
+            hint='Define membership in one or more groups to aid in message routing',
+        )
         # Controller options
         self.add_config_option(
             section='controller',
@@ -194,7 +204,8 @@ def __init__(self, args=None):
             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.',
+            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',
@@ -358,6 +369,11 @@ def parse_options(self, args):
             # because env variables and configparser do not enforce the
             # value type, we do it now to ensure we have the type we want
             self._enforce_entry_type(entry)
+        # Parse plugin_ sections to populate plugin configuration
+        self._config_options['plugins'] = {}
+        if self._config_file:
+            for section in filter(lambda x: x.startswith("plugin_"), self._config_file.sections()):
+                self._config_options['plugins'][section.replace("plugin_", "")] = dict(self._config_file[section])
 
     def _enforce_entry_type(self, entry):
         if entry.value is not None:

receptor/connection.py

@@ -12,8 +12,9 @@
 
 
 class Connection:
-    def __init__(self, id_, protocol_obj, buffer_mgr, receptor):
+    def __init__(self, id_, meta, protocol_obj, buffer_mgr, receptor):
         self.id_ = id_
+        self.meta = meta
         self.protocol_obj = protocol_obj
         self.buffer_mgr = buffer_mgr
         self.receptor = receptor
@@ -53,6 +54,8 @@ def send_route_advertisement(self, edges=None, seen=[]):
                 buf.push(json.dumps({
                     "cmd": "ROUTE",
                     "id": self.receptor.node_id,
+                    "capabilities": self.receptor.work_manager.get_capabilities(),
+                    "groups": self.receptor.config.node_groups,
                     "edges": edges,
                     "seen": seens
                 }).encode("utf-8"))

receptor/messages/directive.py

@@ -1,5 +1,6 @@
 import datetime
 import logging
+import json
 
 from ..exceptions import UnknownDirective
 from . import envelope
@@ -21,7 +22,7 @@ async def __call__(self, router, inner_env):
         if action not in self.CONTROL_DIRECTIVES:
             raise UnknownDirective(f'Unknown control directive: {action}')
         action_method = getattr(self, action)
-        responses = action_method(inner_env)
+        responses = action_method(router.receptor, inner_env)
         serial = 0
         async for response in responses:
             serial += 1
@@ -34,9 +35,12 @@ async def __call__(self, router, inner_env):
             )
             await router.send(enveloped_response)
 
-    async def ping(self, inner_env):
+    async def ping(self, receptor, inner_env):
         logger.info(f'Received ping from {inner_env.sender}')
-        yield f'{inner_env.raw_payload}|{datetime.datetime.utcnow().isoformat()}'
+        return_data = dict(initial_time=inner_env.raw_payload,
+                           response_time=str(datetime.datetime.utcnow()),
+                           active_work=receptor.work_manager.get_work())
+        yield json.dumps(return_data)
 
 
 control = Control()

receptor/protocol.py

@@ -97,11 +97,10 @@ async def wait_greeting(self):
                     self.transport.close()
             await asyncio.sleep(.1)
         logger.debug("handshake complete, starting normal handle loop")
-        self.loop.create_task(self.connection.message_handler(self.incoming_buffer)) # Duplicated (see handle_handshake)?
 
     def handle_handshake(self, data):
         self.greeted = True
-        self.connection = self.receptor.add_connection(data["id"], self)
+        self.connection = self.receptor.add_connection(data["id"], data.get("meta", {}), self)
         self.loop.create_task(self.watch_queue(data["id"], self.transport))
         self.loop.create_task(self.connection.message_handler(self.incoming_buffer))
 
@@ -110,6 +109,9 @@ def send_handshake(self):
             "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())
         }).encode("utf-8")
         self.transport.write(msg + DELIM)
 
@@ -221,4 +223,3 @@ def _data_received_callback(self, inner_env, fut):
                 code=1,
             )
             self.emit_response(err_resp)
-

receptor/receptor.py

@@ -93,9 +93,9 @@ def update_connections(self, connection):
             self.connections[connection.id_] = [connection]
         self.update_connection_manifest(connection.id_)
 
-    def add_connection(self, id_, protocol_obj):
+    def add_connection(self, id_, meta, protocol_obj):
         buffer_mgr = self.config.components_buffer_manager
-        conn = Connection(id_, protocol_obj, buffer_mgr, self)
+        conn = Connection(id_, meta, protocol_obj, buffer_mgr, self)
         self.update_connections(conn)
         return conn
 

receptor/work.py

@@ -10,6 +10,7 @@
 class WorkManager:
     def __init__(self, receptor):
         self.receptor = receptor
+        self.active_work = []
 
     def load_receptor_worker(self, name):
         entry_points = [x for x in filter(lambda x: x.name == name,
@@ -18,6 +19,23 @@ def load_receptor_worker(self, name):
             raise exceptions.UnknownDirective(f"Error loading directive handlers for {name}")
         return entry_points[0].load()
 
+    def get_capabilities(self):
+        return [(x.name, pkg_resources.get_distribution(x.resolve().__package__).version)
+                for x in pkg_resources.iter_entry_points('receptor.worker')]
+
+    def get_work(self):
+        return self.active_work
+
+    def add_work(self, env):
+        self.active_work.append(dict(id=env.message_id,
+                                     directive=env.directive,
+                                     sender=env.sender))
+
+    def remove_work(self, env):
+        for work in self.active_work:
+            if env.message_id == work["id"]:
+                self.active_work.remove(work)
+
     async def handle(self, inner_env):
         logger.info(f'Handling work for {inner_env.message_id} as {inner_env.directive}')
         namespace, action = inner_env.directive.split(':', 1)
@@ -29,8 +47,8 @@ async def handle(self, inner_env):
             except AttributeError:
                 logger.exception(f'Could not load action {action} from {namespace}')
                 raise exceptions.InvalidDirectiveAction(f'Invalid action {action} for {namespace}')
-
-            responses = action_method(inner_env)
+            self.add_work(inner_env)
+            responses = action_method(inner_env, self.receptor.config.plugins.get("namespace", {}))
             async for response in responses:
                 serial += 1
                 logger.debug(f'Response emitted for {inner_env.message_id}, serial {serial}')
@@ -54,5 +72,6 @@ async def handle(self, inner_env):
                 serial=serial,
                 code=1,
             )
+            self.remove_work(inner_env)
             await self.receptor.router.send(enveloped_response)