incoming buffer is now async

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

Author: jhjaggars

receptor/protocol.py

@@ -6,10 +6,8 @@
 import time
 import uuid
 
-from collections import deque
-
-from .messages import envelope
 from .exceptions import ReceptorBufferError
+from .messages import envelope
 from .stats import connected_peers_guage
 
 logger = logging.getLogger(__name__)
@@ -19,19 +17,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):
@@ -74,8 +72,8 @@ 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,18 +89,15 @@ 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

receptor/receptor.py

@@ -1,16 +1,16 @@
-import os
-import json
-import uuid
-import time
 import asyncio
-import logging
 import copy
+import json
+import logging
+import os
+import time
+import uuid
 
+from . import exceptions
+from .messages import directive, envelope
 from .router import MeshRouter
-from .work import WorkManager
-from .messages import envelope, directive
 from .stats import messages_received_counter
-from . import exceptions
+from .work import WorkManager
 
 RECEPTOR_DIRECTIVE_NAMESPACE = 'receptor'
 logger = logging.getLogger(__name__)
@@ -99,12 +99,11 @@ def update_connections(self, protocol_obj):
 
     async def message_handler(self, buf):
         while True:
-            for data in buf.get():
-                if "cmd" in data and data["cmd"] == "ROUTE":
-                    self.handle_route_advertisement(data)
-                else:
-                    await self.handle_message(data)
-            await asyncio.sleep(.1)
+            data = await buf.get()
+            if "cmd" in data and data["cmd"] == "ROUTE":
+                self.handle_route_advertisement(data)
+            else:
+                await self.handle_message(data)
 
     def add_connection(self, protocol_obj):
         self.update_connections(protocol_obj)

receptor/tests/test_protocol.py

@@ -1,35 +1,34 @@
 import pytest
+
 from receptor import protocol
 
 
 def deser(x):
     return x
 
 
-def test_databuffer():
+@pytest.mark.asyncio
+async def test_databuffer():
     b = protocol.DataBuffer(deserializer=deser)
     msg = b"this is a test"
     s = protocol.DELIM + msg + protocol.DELIM
     b.add(s)
-    it = b.get()
-    assert next(it) == b""
-    assert next(it) == msg
+    assert await b.get() == b""
+    assert await b.get() == msg
 
 
 def test_databuffer_no_delim():
     b = protocol.DataBuffer(deserializer=deser)
     msg = b"this is a test"
     b.add(msg)
-    with pytest.raises(StopIteration):
-        next(b.get())
+    assert b.q.empty()
 
 
-def test_databuffer_many_msgs():
+@pytest.mark.asyncio
+async def test_databuffer_many_msgs():
     b = protocol.DataBuffer(deserializer=deser)
     msg = [b"first bit", b"second bit", b"third bit unfinished"]
     b.add(protocol.DELIM.join(msg))
-    it = b.get()
-    assert next(it) == msg[0]
-    assert next(it) == msg[1]
-    with pytest.raises(StopIteration):
-        next(it)
+    assert msg[0] == await b.get()
+    assert msg[1] == await b.get()
+    assert b.q.empty()