Implement message expiration at the connection level

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

@@ -28,6 +28,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 +49,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 +88,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/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

@@ -34,9 +34,6 @@ def __init__(self, receptor, loop):
         self.receptor = receptor
         self.loop = loop
 
-    async def watch_message_expire(self):
-        pass
-
     async def watch_queue(self, node, transport):
         buffer_mgr = self.receptor.config.components.buffer_manager
         buffer_obj = buffer_mgr.get_buffer_for_node(node, self.receptor)

receptor/receptor.py

@@ -4,6 +4,7 @@
 import time
 import asyncio
 import logging
+import copy
 
 from .config import ReceptorConfig
 from .router import MeshRouter
@@ -37,6 +38,26 @@ 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:
+            logger.info("Checking expirations")
+            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)
+                    logger.info("Examining {}".format(message_data))
+                    if "expire_time" in message_data and message_data['expire_time'] < time.time():
+                        logger.info("Expiring message {}:{}".format(ident, connection["id"]))
+                        expired_message = buffer.read_message(ident, remove=True)
+                        # 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 []
@@ -48,43 +69,46 @@ def get_connection_manifest(self):
             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.id_:
+            if node["id"] == connection:
                 node["last"] = time.time()
                 found = True
                 break
         if not found:
-            node.append(dict(id=connection.id_,
-                             last=time.time()))
-        fd = open(self.connection_manifest_path, "w")
-        json.dump(manifest, fd)
-        fd.close()
+            manifest.append(dict(id=connection.id_,
+                            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)
+        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):
+    def remove_connection(self, protocol_obj):
         notify_protocols = []
-        self.update_connection_manifest(conn)
+        #self.update_connection_manifest(conn)
         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]