wiring DurableBuffer up

Signed-off-by: Jesse Jaggars <jjaggars@redhat.com>

Author: jhjaggars

receptor/buffers/file.py

@@ -6,8 +6,7 @@
 from collections import deque
 from concurrent.futures import ThreadPoolExecutor
 
-from ..exceptions import ReceptorBufferError
-from .base import BaseBuffer, BaseBufferManager
+from .base import BaseBufferManager
 
 logger = logging.getLogger(__name__)
 pool = ThreadPoolExecutor()
@@ -33,24 +32,23 @@ def __init__(self, dir_, key, loop):
         self.q._queue = deque(self._read_manifest())
 
     async def put(self, data):
-        async with self._manifest_lock:
-            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()
+        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):
-        async with self._manifest_lock:
-            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
+        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):
-        await self._loop.run_in_executor(pool, self._write_manifest)
+        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:
@@ -90,91 +88,5 @@ 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.loop = asyncio.get_event_loop()
-        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
-
-    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):
-        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 []
-        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))
-            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 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))
-        self.write_manifest(manifest)
-        return item
+        path = os.path.join(os.path.expanduser(receptor.config.default_data_dir))
+        return self._buffers.setdefault(node_id, DurableBuffer(path, node_id, asyncio.get_event_loop()))

receptor/buffers/memory.py

@@ -6,7 +6,6 @@
 import time
 import uuid
 
-from .exceptions import ReceptorBufferError
 from .messages import envelope
 from .stats import connected_peers_guage
 
@@ -52,19 +51,13 @@ 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()
+                logger.debug("about to wait for a message")
+                msg = await buffer_obj.get()
+                logger.debug("got a message to send")
                 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
 
@@ -126,7 +119,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):
@@ -156,7 +149,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):

receptor/protocol.py

@@ -101,7 +101,7 @@ async def message_handler(self, buf):
         while True:
             data = await buf.get()
             if "cmd" in data and data["cmd"] == "ROUTE":
-                self.handle_route_advertisement(data)
+                await self.handle_route_advertisement(data)
             else:
                 await self.handle_message(data)
 
@@ -119,19 +119,19 @@ def remove_connection(self, protocol_obj):
                 self.router.debug_router()
                 self.update_connection_manifest(connection_node)
             notify_connections += self.connections[connection_node]
-        self.send_route_advertisement(self.router.get_edges())
+        protocol_obj.loop.create_task(self.send_route_advertisement(self.router.get_edges()))
 
     async def shutdown_handler(self):
         while True:
             if self.stop:
                 return
             await asyncio.sleep(1)
 
-    def handle_route_advertisement(self, data):
+    async def handle_route_advertisement(self, data):
         self.router.add_edges(data["edges"])
-        self.send_route_advertisement(data["edges"], data["seen"])
+        await self.send_route_advertisement(data["edges"], data["seen"])
 
-    def send_route_advertisement(self, edges=None, seen=[]):
+    async def send_route_advertisement(self, edges=None, seen=[]):
         edges = edges or self.router.get_edges()
         seen = set(seen)
         logger.debug("Emitting Route Advertisements, excluding {}".format(seen))
@@ -142,17 +142,14 @@ def send_route_advertisement(self, edges=None, seen=[]):
         for target in destinations:
             buf = self.buffer_mgr.get_buffer_for_node(target, self)
             try:
-                buf.push(json.dumps({
+                await buf.put(json.dumps({
                     "cmd": "ROUTE",
                     "id": self.node_id,
                     "capabilities": self.work_manager.get_capabilities(),
                     "groups": self.config.node_groups,
                     "edges": edges,
                     "seen": seens
                 }).encode("utf-8"))
-            except exceptions.ReceptorBufferError as e:
-                logger.exception("Receptor Buffer Write Error broadcasting routes and capabilities: {}".format(e))
-                # TODO: This might should be a hard shutdown event
             except Exception as e:
                 logger.exception("Error trying to broadcast routes and capabilities: {}".format(e))
 

receptor/receptor.py

@@ -1,14 +1,14 @@
-import logging
-
 import datetime
-from collections import defaultdict
 import heapq
+import logging
 import random
 import uuid
+from collections import defaultdict
 
 from dateutil import parser
+
+from .exceptions import ReceptorBufferError, UnrouteableError
 from .messages import envelope
-from .exceptions import UnrouteableError, ReceptorBufferError
 from .stats import route_counter
 
 logger = logging.getLogger(__name__)
@@ -147,7 +147,7 @@ async def forward(self, outer_envelope, next_hop):
         logger.debug(f'Forwarding frame {outer_envelope.frame_id} to {next_hop}')
         try:
             route_counter.inc()
-            buffer_obj.push(outer_envelope.serialize().encode("utf-8"))
+            await buffer_obj.put(outer_envelope.serialize().encode("utf-8"))
         except ReceptorBufferError as e:
             logger.exception("Receptor Buffer Write Error forwarding message to {}: {}".format(next_hop, e))
             # TODO: Possible to find another route? This might be a hard failure
@@ -188,6 +188,3 @@ async def send(self, inner_envelope, expected_response=False):
         if expected_response and inner_envelope.message_type == 'directive':
             self.response_registry[inner_envelope.message_id] = dict(message_sent_time=inner_envelope.timestamp)
         await self.forward(outer_envelope, next_node_id)
-
-
-