Merge pull request #71 from jhjaggars/fix-buffer-expire

storing expire time as metadata rather than part of the payload

Author: matburt

receptor/buffers/file.py

@@ -1,8 +1,8 @@
 import asyncio
+import datetime
 import json
 import logging
 import os
-import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 
@@ -12,6 +12,29 @@
 pool = ThreadPoolExecutor()
 
 
+class ManifestEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, datetime.datetime):
+            return {
+                "_type": "datetime.datetime",
+                "value": o.isoformat(),
+            }
+        return super().default(o)
+
+
+class ManifestDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(object_hook=self.object_hook, *args, **kwargs)
+
+    def object_hook(self, o):
+        type_ = o.get("_type")
+        if type_ != "datetime.datetime":
+            return o
+
+        t = o['_type']
+        return datetime.datetime.fromisoformat(o["value"])
+
+
 class DurableBuffer:
 
     def __init__(self, dir_, key, loop):
@@ -29,17 +52,20 @@ def __init__(self, dir_, key, loop):
             self.q.put_nowait(item)
 
     async def put(self, data):
-        ident = str(uuid.uuid4())
-        await self._loop.run_in_executor(pool, self._write_file, data, ident)
-        await self.q.put(ident)
+        item = {
+            "ident": str(uuid.uuid4()),
+            "expire_time": datetime.datetime.utcnow() + datetime.timedelta(minutes=5),
+        }
+        await self._loop.run_in_executor(pool, self._write_file, data, item)
+        await self.q.put(item)
         await self._save_manifest()
 
     async def get(self, handle_only=False, delete=True):
         while True:
             msg = await self.q.get()
             await self._save_manifest()
             try:
-                return await self._get_file(msg, handle_only=handle_only, delete=delete)
+                return await self._get_file(msg["ident"], handle_only=handle_only, delete=delete)
             except FileNotFoundError:
                 pass
 
@@ -49,12 +75,12 @@ async def _save_manifest(self):
 
     def _write_manifest(self):
         with open(self._manifest_path, "w") as fp:
-            json.dump(list(self.q._queue), fp)
+            json.dump(list(self.q._queue), fp, cls=ManifestEncoder)
 
     def _read_manifest(self):
         try:
             with open(self._manifest_path, "r") as fp:
-                return json.load(fp)
+                return json.load(fp, cls=ManifestDecoder)
         except FileNotFoundError:
             return []
         except json.decoder.JSONDecodeError:
@@ -77,30 +103,29 @@ async def _get_file(self, ident, handle_only=False, delete=True):
         fp = await self._loop.run_in_executor(pool, open, path, "rb")
         if handle_only:
             return fp
-        bytes = await self._loop.run_in_executor(pool, lambda: fp.read())
+        bytes_ = await self._loop.run_in_executor(pool, fp.read)
         fp.close()
         if delete:
             await self._loop.run_in_executor(pool, os.remove, path)
-        return bytes
+        return bytes_
 
-    def _write_file(self, data, ident):
-        with open(os.path.join(self._message_path, ident), "wb") as fp:
+    def _write_file(self, data, item):
+        with open(os.path.join(self._message_path, item["ident"]), "wb") as fp:
             fp.write(data)
 
     async def expire(self):
         async with self._manifest_lock:
             new_queue = asyncio.Queue()
             while self.q.qsize() > 0:
-                ident = await self.q.get()
-                data = await self._get_file(ident, handle_only=True, delete=False)
-                # TODO: This will never work, it's not pure json anymore
-                msg = json.load(data)
-                if "expire_time" in msg and msg['expire_time'] < time.time():
+                item = await self.q.get()
+                ident = item["ident"]
+                expire_time = item["expire_time"]
+                if expire_time < datetime.datetime.utcnow():
                     logger.info("Expiring message %s", ident)
                     # TODO: Do something with expired message
                     await self._loop.run_in_executor(pool, os.remove, self._path_for_ident(ident))
                 else:
-                    await new_queue.put(ident)
+                    await new_queue.put(item)
             self.q = new_queue
 
 

receptor/receptor.py

@@ -1,4 +1,3 @@
-import pkg_resources
 import asyncio
 import copy
 import json
@@ -7,6 +6,8 @@
 import time
 import uuid
 
+import pkg_resources
+
 from . import exceptions
 from .messages import directive, envelope
 from .router import MeshRouter