event driven approach

lib/python/flame/backend/chunk_store.py

@@ -13,11 +13,14 @@
 # limitations under the License.
 #
 # SPDX-License-Identifier: Apache-2.0
-
+"""ChunkStore."""
 
 import logging
+from threading import Thread
 from typing import Tuple, Union
 
+from ..common.constants import EMPTY_PAYLOAD
+from ..common.util import run_async
 from ..proto import backend_msg_pb2 as msg_pb2
 
 DEFAULT_CHUNK_SIZE = 1048576  # 1MB
@@ -26,17 +29,28 @@
 
 
 class ChunkStore(object):
-    def __init__(self):
+    """ChunkStore class."""
+
+    def __init__(self, loop=None, channel=None):
+        """Initialize an instance."""
+        self._loop = loop
+        self._channel = channel
+
         # for fragment
         self.pidx = 0
         self.cidx = DEFAULT_CHUNK_SIZE
 
+        # for assemble
+        self.recv_buf = list()
+
         # for both fragment and assemble
         self.data = b''
         self.seqno = -1
         self.eom = False
 
     def set_data(self, data: bytes) -> None:
+        """Set data in chunk store."""
+        logger.debug(f"setting data of size {len(data)}")
         self.data = data
 
         # reset variables since new data is set
@@ -45,8 +59,10 @@ def set_data(self, data: bytes) -> None:
 
     def get_chunk(self) -> Tuple[Union[bytes, None], int, bool]:
         """
-        get_chunk() returns None as the first part of the triplet
-        if its internal index is pointing beyond the end of message.
+        Return a chunk of data.
+
+        The method returns None as the first part of the triplet
+        if its internal index is pointing beyond the end of data.
         Otherwise, it returns a chunk every time it is called.
         """
         data_len = len(self.data)
@@ -67,16 +83,54 @@ def get_chunk(self) -> Tuple[Union[bytes, None], int, bool]:
         self.pidx = self.cidx
         self.cidx += DEFAULT_CHUNK_SIZE
 
+        logger.debug(f"chunk {seqno}: {len(data)}")
         return data, seqno, eom
 
     def assemble(self, msg: msg_pb2.Data) -> bool:
+        """Assemble message.
+
+        This method pushes message payload into a receive buffer.
+        If eom (end of message) is set, bytes in the array are joined.
+        Then, the assembled data will be put into a receive queue.
+
+        The join operation can be exepnsive if the data size is large.
+        We run the join operation in a separate thread in order to unblock
+        asyncio tasks as quickly as possible.
+        """
         # out of order delivery
         if self.seqno + 1 != msg.seqno:
             logger.warning(f'out-of-order seqno from {msg.end_id}')
             return False
 
-        self.data += msg.payload
+        logger.debug(f"chunk {msg.seqno}: {len(msg.payload)}")
+        # add payload to a recv buf
+        self.recv_buf.append(msg.payload)
         self.seqno = msg.seqno
         self.eom = msg.eom
 
+        if self.eom:
+            # we assemble payloads in the recv buf array
+            # only if eom is set to True.
+            # In this way, we only pay byte concatenation cost once
+            _thread = Thread(target=self._assemble,
+                             args=(msg.end_id, ),
+                             daemon=True)
+            _thread.start()
+
         return True
+
+    def _assemble(self, end_id: str) -> None:
+        # This code must be executed in a separate thread
+        self.data = EMPTY_PAYLOAD.join(self.recv_buf)
+
+        async def inner():
+            logger.debug(f'fully assembled data size = {len(self.data)}')
+
+            rxq = self._channel.get_rxq(end_id)
+            if rxq is None:
+                logger.debug(f"rxq not found for {end_id}")
+                return
+
+            await rxq.put(self.data)
+
+        _, status = run_async(inner(), self._loop)

lib/python/flame/backend/mqtt.py

@@ -20,7 +20,6 @@
 import time
 from collections import deque
 from enum import IntEnum
-from typing import Tuple
 
 import paho.mqtt.client as mqtt
 from google.protobuf.any_pb2 import Any
@@ -247,18 +246,14 @@ async def _handle_data(self, any_msg: Any) -> None:
             expiry = time.time() + MQTT_TIME_WAIT
             self._cleanup_waits[msg.end_id] = expiry
 
-        payload, fully_assembled = self.assemble_chunks(msg)
-
-        channel = self._channels[msg.channel_name]
-
-        if fully_assembled:
-            logger.debug(f'fully assembled data size = {len(payload)}')
-            rxq = channel.get_rxq(msg.end_id)
-            if rxq is None:
-                logger.debug(f"rxq not found for {msg.end_id}")
-                return
+        if msg.end_id not in self._msg_chunks:
+            channel = self._channels[msg.channel_name]
+            self._msg_chunks[msg.end_id] = ChunkStore(self._loop, channel)
 
-            await rxq.put(payload)
+        chunk_store = self._msg_chunks[msg.end_id]
+        if not chunk_store.assemble(msg) or chunk_store.eom:
+            # clean up if message is wrong or completely assembled
+            del self._msg_chunks[msg.end_id]
 
     async def _rx_task(self):
         self._rx_deque = deque()
@@ -315,24 +310,6 @@ def on_message(self, client, userdata, message):
         # add one extra future in the queue
         self._rx_deque.append(self._loop.create_future())
 
-    def assemble_chunks(self, msg: msg_pb2.Data) -> Tuple[bytes, bool]:
-        """Assemble message chunks to build a whole message."""
-        if msg.end_id not in self._msg_chunks:
-            self._msg_chunks[msg.end_id] = ChunkStore()
-
-        chunk_store = self._msg_chunks[msg.end_id]
-        if not chunk_store.assemble(msg):
-            # clean up wrong message
-            del self._msg_chunks[msg.end_id]
-            return b'', False
-
-        if chunk_store.eom:
-            data = chunk_store.data
-            del self._msg_chunks[msg.end_id]
-            return data, True
-        else:
-            return b'', False
-
     def subscribe(self, topic) -> None:
         """Subscribe to a topic."""
         logger.debug(f'subscribe topic: {topic}')