Merge pull request #40 from jhjaggars/async-buffers

Adding Async buffers

Author: matburt

receptor/buffers/file.py

@@ -1,102 +1,113 @@
-import logging
-import uuid
+import asyncio
 import json
+import logging
 import os
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
 
-from .base import BaseBufferManager, BaseBuffer
-from ..exceptions import ReceptorBufferError
+from .base import BaseBufferManager
 
 logger = logging.getLogger(__name__)
+pool = ThreadPoolExecutor()
 
 
-class FileBufferManager(BaseBufferManager):
-    _buffers = {}
-
-    def get_buffer_for_node(self, node_id, receptor):
-        return self._buffers.setdefault(node_id, FileBuffer(node_id, receptor))
-
-
-class FileBuffer(BaseBuffer):
-
-    def __init__(self, node_id, receptor):
-        super().__init__(node_id, receptor)
-        self.node_id = node_id
-        self.base_path = os.path.join(os.path.expanduser(self.receptor.config.default_data_dir))
-        self.message_path = os.path.join(self.base_path, self.receptor.node_id, "messages")
-        self.manifest_path = os.path.join(self.base_path, self.receptor.node_id, "manifest-{}".format(node_id))
-        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
+class DurableBuffer:
 
-    def new_message(self):
-        ident = str(uuid.uuid4())
-        try:
-            handle = open(os.path.join(self.message_path, ident), "wb")
-        except Exception as e:
-            raise ReceptorBufferError("Failed to generate new message file for {}: {}".format(self.node_id, e))
-        return (ident, handle)
-
-    def read_message(self, ident, remove=True):
+    def __init__(self, dir_, key, loop):
+        self.q = asyncio.Queue()
+        self._base_path = os.path.join(os.path.expanduser(dir_))
+        self._message_path = os.path.join(self._base_path, "messages")
+        self._manifest_path = os.path.join(self._base_path, f"manifest-{key}")
+        self._loop = loop
+        self._manifest_lock = asyncio.Lock(loop=self._loop)
         try:
-            message_data = open(os.path.join(self.message_path, ident), "rb").read()
-            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
-
-    def write_manifest(self, manifest):
-        try:
-            fd = open(self.manifest_path, "w")
-            json.dump(manifest, fd)
-            fd.close()
-        except Exception as e:
-            raise ReceptorBufferError("Failed to handle metadata file for {}: {}".format(self.node_id, e))
-
-    def read_manifest(self):
-        if not os.path.exists(self.manifest_path):
-            return []
+            os.makedirs(self._message_path, mode=0o700)
+        except Exception:
+            pass
+        for item in self._read_manifest():
+            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)
+        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)
+            except FileNotFoundError:
+                pass
+    
+    async def _save_manifest(self):
+        async with self._manifest_lock:
+            await self._loop.run_in_executor(pool, self._write_manifest)
+    
+    def _write_manifest(self):
+        with open(self._manifest_path, "w") as fp:
+            json.dump(list(self.q._queue), fp)
+    
+    def _read_manifest(self):
         try:
-            fd = open(self.manifest_path, "r")
-            manifest = json.load(fd)
-            return manifest
-        except Exception as e:
-            logger.warn("Failed to read manifest: {}".format(e))
+            with open(self._manifest_path, "r") as fp:
+                return json.load(fp)
+        except FileNotFoundError:
             return []
 
-    def push(self, message):
-        manifest = self.read_manifest()
-        ident, handle = self.new_message()
-        try:
-            handle.write(message)
-            handle.close()
-        except Exception as e:
-            raise ReceptorBufferError("Failed to write message file for {} {}: {}".format(self.node_id, ident, e))
-        manifest.append(ident)
-        self.write_manifest(manifest)
+    def _path_for_ident(self, ident):
+        return os.path.join(self._message_path, ident)
+
+    async def _get_file(self, ident, handle_only=False, delete=True):
+        """
+        Retrieves a file from disk. If handle_only is True then we will
+        return the handle to the file and do nothing else. Otherwise the file
+        is read into memory all at once and returned. If delete is True (the
+        default) and handle_only is False (the default) then the underlying
+        file will be removed as well.
+        """
+        path = self._path_for_ident(ident)
+        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())
+        fp.close()
+        if delete:
+            await self._loop.run_in_executor(pool, os.remove, path)
+        return bytes
+
+    def _write_file(self, data, ident):
+        with open(os.path.join(self._message_path, 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)
+                msg = json.load(data)
+                if "expire_time" in msg and msg['expire_time'] < time.time():
+                    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)
+            self.q = new_queue
+
 
-    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
+class FileBufferManager(BaseBufferManager):
+    _buffers = {}
 
-    def pop(self):
-        manifest = self.read_manifest()
-        item = self.read_message(manifest.pop(0))
-        self.write_manifest(manifest)
-        return item
+    def get_buffer_for_node(self, node_id, receptor):
+        # due to the way that the manager is constructed, we won't have enough
+        # information to build a proper defaultdict at the time, and we want to
+        # make sure we only construct a single instance of DurableBuffer
+        # per-node so.. doing this the hard way.
+        if node_id not in self._buffers:
+            path = os.path.join(os.path.expanduser(receptor.config.default_data_dir))
+            self._buffers[node_id] = DurableBuffer(path, node_id, asyncio.get_event_loop())
+        return self._buffers[node_id]

receptor/buffers/memory.py

@@ -6,10 +6,7 @@
 import time
 import uuid
 
-from collections import deque
-
 from .messages import envelope
-from .exceptions import ReceptorBufferError
 from .stats import connected_peers_guage
 
 logger = logging.getLogger(__name__)
@@ -19,19 +16,19 @@
 
 
 class DataBuffer:
-    def __init__(self, deserializer=json.loads):
-        self.q = deque()
+    def __init__(self, loop=None, deserializer=json.loads):
+        self.q = asyncio.Queue(loop=loop)
         self.data_buffer = b""
         self.deserializer = deserializer
 
     def add(self, data):
         self.data_buffer = self.data_buffer + data
         *ready, self.data_buffer = self.data_buffer.rsplit(DELIM)
-        self.q.extend(ready)
+        for chunk in ready:
+            self.q.put_nowait(chunk)
 
-    def get(self):
-        while self.q:
-            yield self.deserializer(self.q.popleft())
+    async def get(self):
+        return self.deserializer(await self.q.get())
 
 
 class BaseProtocol(asyncio.Protocol):
@@ -54,28 +51,24 @@ async def watch_queue(self):
         buffer_obj = buffer_mgr.get_buffer_for_node(self.id, self.receptor)
         while not self.transport.is_closing():
             try:
-                msg = buffer_obj.pop()
+                msg = await buffer_obj.get()
+            except Exception:
+                logger.exception("Unhandled error when fetch from buffer for %s", self.id)
+                continue
+
+            try:
                 self.transport.write(msg + DELIM)
-            except IndexError:
-                await asyncio.sleep(0.1)
-            except ReceptorBufferError as e:
-                logger.exception("Receptor Buffer Read Error: {}".format(e))
-                # TODO: We need to try to send this message along somewhere else
-                # and record the failure somewhere
-                self.transport.close()
-                return
-            except Exception as e:
-                logger.exception("Error received trying to write to {}: {}".format(self.id, e))
-                buffer_obj.push(msg)
+            except Exception:
+                logger.exception("Error received trying to write to %s", self.id)
+                await buffer_obj.put(msg)
                 self.transport.close()
                 return
 
     def connection_made(self, transport):
         self.peername = transport.get_extra_info('peername')
         self.transport = transport
-        self.greeted = False
-        self.incoming_buffer = DataBuffer()
         connected_peers_guage.inc()
+        self.incoming_buffer = DataBuffer(loop=self.loop)
         self.loop.create_task(self.wait_greeting())
 
     def connection_lost(self, exc):
@@ -91,21 +84,17 @@ 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():
-                logger.debug(data)
-                if data["cmd"] == "HI":
-                    self.handle_handshake(data)
-                    break
-                else:
-                    logger.error("Handshake failed!")
-                    self.transport.close()
-            await asyncio.sleep(.1)
-        logger.debug("handshake complete, starting normal handle loop")
+        logger.debug('Looking for handshake...')
+        data = await self.incoming_buffer.get()
+        logger.debug(data)
+        if data["cmd"] == "HI":
+            self.handle_handshake(data)
+            logger.debug("handshake complete, starting normal handle loop")
+        else:
+            logger.error("Handshake failed!")
+            self.transport.close()
 
     def handle_handshake(self, data):
-        self.greeted = True
         self.id = data["id"]
         self.meta = data.get("meta", {})
         self.receptor.add_connection(self)
@@ -133,7 +122,7 @@ def handle_handshake(self, data):
         super().handle_handshake(data)
         logger.debug("Received handshake from client with id %s, responding...", data["id"])
         self.send_handshake()
-        self.receptor.send_route_advertisement()
+        self.loop.create_task(self.receptor.send_route_advertisement())
 
 
 async def create_peer(receptor, loop, host, port):
@@ -163,7 +152,7 @@ def connection_lost(self, exc):
     def handle_handshake(self, data):
         super().handle_handshake(data)
         logger.debug("Received handshake from server with id %s", data["id"])
-        self.receptor.send_route_advertisement()
+        self.loop.create_task(self.receptor.send_route_advertisement())
 
 
 class BasicControllerProtocol(asyncio.Protocol):