Merge pull request #20 from matburt/message_expiry

Add envelop element for expiration, track connections persistently

Author: matburt

receptor/buffers/base.py

@@ -14,6 +14,12 @@ def __init__(self, node_id, receptor):
         self.node_id = node_id
         self.receptor = receptor
 
+    def __iter__(self):
+        raise NotImplementedError
+
+    def __next__(self):
+        raise NotImplementedError
+
     def push(self, message):
         raise NotImplementedError()
 

receptor/buffers/file.py

@@ -1,4 +1,3 @@
-import collections
 import logging
 import uuid
 import json
@@ -28,6 +27,18 @@ def __init__(self, node_id, receptor):
         if not os.path.exists(self.message_path):
             os.makedirs(self.message_path, mode=0o700)
 
+    def __iter__(self):
+        self.current = 0
+        return self
+
+    def __next__(self):
+        manifest = self.read_manifest()
+        if len(manifest) <= self.current:
+            raise StopIteration
+        ident = manifest[self.current]
+        current_payload = self.read(ident)
+        self.current += 1
+        return ident, current_payload
 
     def new_message(self):
         ident = str(uuid.uuid4())
@@ -37,10 +48,11 @@ def new_message(self):
             raise ReceptorBufferError("Failed to generate new message file for {}: {}".format(self.node_id, e))
         return (ident, handle)
 
-    def read_message(self, ident):
+    def read_message(self, ident, remove=True):
         try:
             message_data = open(os.path.join(self.message_path, ident), "rb").read()
-            os.remove(os.path.join(self.message_path, ident))
+            if remove:
+                os.remove(os.path.join(self.message_path, ident))
         except Exception as e:
             raise ReceptorBufferError("Failed to handle message data file for {} {}: {}".format(self.node_id, ident, e))
         return message_data
@@ -75,6 +87,14 @@ def push(self, message):
         manifest.append(ident)
         self.write_manifest(manifest)
 
+    def read(self, ident, remove=False):
+        manifest = self.read_manifest()
+        message = self.read_message(ident, remove=remove)
+        if remove:
+            manifest.remove(ident)
+            self.write_manifest(manifest)
+        return message
+
     def pop(self):
         manifest = self.read_manifest()
         item = self.read_message(manifest.pop(0))

receptor/controller.py

@@ -35,6 +35,7 @@ def mainloop(receptor, socket_path, loop=asyncio.get_event_loop()):
     )
     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:

receptor/messages/envelope.py

@@ -3,6 +3,7 @@
 import json
 import logging
 import uuid
+import time
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +38,7 @@ def serialize(self):
 class InnerEnvelope:
     def __init__(self, receptor, message_id, sender, recipient, message_type, timestamp,
                  raw_payload, directive=None, in_response_to=None, ttl=None,
-                 serial=1):
+                 serial=1, expire_time_delta=300):
         self.receptor = receptor
         self.message_id = message_id
         self.sender = sender
@@ -48,6 +49,9 @@ def __init__(self, receptor, message_id, sender, recipient, message_type, timest
         self.directive = directive # None if response, 'namespace:action' if not
         self.in_response_to = in_response_to # None if directive, a message_id if not
         self.ttl = ttl # Optional
+        if not expire_time_delta:
+            self.expire_time = None
+        self.expire_time = time.time() + expire_time_delta
         self.serial = serial # serial index of responses
 
     @classmethod

receptor/node.py

@@ -19,6 +19,7 @@ def mainloop(receptor, ping_interval=None, loop=asyncio.get_event_loop(), skip_r
     if ping_interval:
         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))
+    loop.create_task(receptor.watch_expire())
     if not skip_run:
         try:
             loop.run_until_complete(receptor.shutdown_handler())

receptor/protocol.py

@@ -1,6 +1,7 @@
 import datetime
 import asyncio
 import logging
+import time
 import json
 import uuid
 from collections import deque
@@ -35,6 +36,11 @@ def __init__(self, receptor, loop):
         self.loop = loop
 
     async def watch_queue(self, node, transport):
+        '''
+        Watches the buffer for this connection for messages delivered from other
+        parts of Receptor (forwarded messages for example) for messages to send
+        over the connection.
+        '''
         buffer_mgr = self.receptor.config.components.buffer_manager
         buffer_obj = buffer_mgr.get_buffer_for_node(node, self.receptor)
         while not transport.is_closing():
@@ -63,13 +69,17 @@ def connection_made(self, transport):
         self.loop.create_task(self.wait_greeting())
 
     def connection_lost(self, exc):
-        self.receptor.remove_connection(self)
+        self.receptor.remove_connection(self.connection)
 
     def data_received(self, data):
         logger.debug(data)
         self.incoming_buffer.add(data)
 
     async def wait_greeting(self):
+        '''
+        Initialized when the connection is established to handle the greeting
+        before transitioning to message processing.
+        '''
         while not self.greeted:
             logger.debug('Looking for handshake...')
             for data in self.incoming_buffer.get():
@@ -79,10 +89,10 @@ async def wait_greeting(self):
                     break
                 else:
                     logger.error("Handshake failed!")
-                    # TODO: Trigger disconnection
+                    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?
+        self.loop.create_task(self.connection.message_handler(self.incoming_buffer)) # Duplicated (see handle_handshake)?
 
     def handle_handshake(self, data):
         self.greeted = True
@@ -94,6 +104,7 @@ def send_handshake(self):
         msg = json.dumps({
             "cmd": "HI",
             "id": self.receptor.node_id,
+            "expire_time": time.time() + 10,
         }).encode("utf-8")
         self.transport.write(msg + DELIM)
 
@@ -175,7 +186,6 @@ def data_received(self, data):
             raw_payload=payload,
             directive=directive
         )
-        # TODO: Response expiration task?
         # TODO: Persistent registry?
         self.loop.create_task(self.receptor.router.send(inner_env,
                                                         expected_response=True))

receptor/receptor.py

@@ -1,7 +1,10 @@
 import os
+import json
 import uuid
+import time
 import asyncio
 import logging
+import copy
 
 from .config import ReceptorConfig
 from .router import MeshRouter
@@ -20,6 +23,9 @@ def __init__(self, config=None, node_id=None, router_cls=None,
         self.work_manager = (work_manager_cls or WorkManager)(self)
         self.connections = dict()
         self.controller_connections = []
+        self.connection_manifest_path = os.path.join(self.config.server.data_dir,
+                                                     self.node_id,
+                                                     "connection_manifest")
         self.stop = False
 
     def _find_node_id(self):
@@ -32,31 +38,80 @@ def _find_node_id(self):
                 ofs.write(f'\nRECEPTOR_NODE_ID={node_id}\n')
         return str(node_id)
 
+    async def watch_expire(self):
+        while True:
+            current_manifest = self.get_connection_manifest()
+            for connection in current_manifest:
+                buffer = self.config.components.buffer_manager.get_buffer_for_node(connection["id"], self)
+                for ident, message in buffer:
+                    message_actual = json.loads(message)
+                    if "expire_time" in message_actual and message_actual['expire_time'] < time.time():
+                        buffer.read_message(ident, remove=True)
+                        logger.info("Expiring message {}:{}".format(ident, connection["id"]))
+                        # TODO: Do something with expired message
+                if connection["last"] + 86400 < time.time():
+                    logger.info("Expiring connection {}".format(connection["id"]))
+                    write_manifest = copy.copy(current_manifest)
+                    write_manifest.remove(connection)
+                    self.write_connection_manifest(write_manifest)
+            await asyncio.sleep(600)
+
+    def get_connection_manifest(self):
+        if not os.path.exists(self.connection_manifest_path):
+            return []
+        try:
+            fd = open(self.connection_manifest_path, "r")
+            manifest = json.load(fd)
+            return manifest
+        except Exception as e:
+            logger.warn("Failed to read connection manifest: {}".format(e))
+            return []
+
+    def write_connection_manifest(self, manifest):
+        fd = open(self.connection_manifest_path, "w")
+        json.dump(manifest, fd)
+        fd.close()
+
+    def update_connection_manifest(self, connection):
+        manifest = self.get_connection_manifest()
+        found = False
+        for node in manifest:
+            if node["id"] == connection:
+                node["last"] = time.time()
+                found = True
+                break
+        if not found:
+            manifest.append(dict(id=connection,
+                            last=time.time()))
+        self.write_connection_manifest(manifest)
+                        
     def update_connections(self, connection):
         self.router.register_edge(connection.id_, self.node_id, 1)
         if connection.id_ in self.connections:
             self.connections[connection.id_].append(connection)
         else:
             self.connections[connection.id_] = [connection]
+        self.update_connection_manifest(connection.id_)
 
     def add_connection(self, id_, protocol_obj):
         buffer_mgr = self.config.components.buffer_manager
         conn = Connection(id_, protocol_obj, buffer_mgr, self)
         self.update_connections(conn)
         return conn
 
-    def remove_connection(self, conn):
-        notify_protocols = []
+    def remove_connection(self, protocol_obj):
+        notify_connections = []
         for connection_node in self.connections:
-            if conn in self.connections[connection_node]:
-                logger.info("Removing connection {} for node {}".format(conn, connection_node))
-                self.connections[connection_node].remove(conn)
+            if protocol_obj in self.connections[connection_node]:
+                logger.info("Removing connection {} for node {}".format(protocol_obj, connection_node))
+                self.update_connection_manifest(connection_node)
+                self.connections[connection_node].remove(protocol_obj)
                 self.router.update_node(self.node_id, connection_node, 100)
                 self.router.debug_router()
-            notify_protocols += self.connections[connection_node]
-        # TODO: Broadcast update, set timer for full expiration
-        for active_protocol in notify_protocols:
-            active_protocol.send_route_advertisement(self.router.get_edges())
+                self.update_connection_manifest(connection_node)
+            notify_connections += self.connections[connection_node]
+        for active_connection in notify_connections:
+            active_connection.send_route_advertisement(self.router.get_edges())
 
     async def shutdown_handler(self):
         while True: