Added Batch signal subscriber. (#9)

* Added Batch signal subscriber.

* Import sorting

Author: DustinMoriarty

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "asyncio-signal-bus"
-version = "1.2.1"
+version = "1.3.0"
 description = "Internal application publisher/subscriber bus using asyncio queues."
 authors = ["DustinMoriarty <dustin.moriarty@protonmail.com>"]
 readme = "README.md"

src/asyncio_signal_bus/signal_bus.py

@@ -7,7 +7,7 @@
 from asyncio_signal_bus.injector import Injector
 from asyncio_signal_bus.publisher import SignalPublisher
 from asyncio_signal_bus.queue_getter import QueueGetter
-from asyncio_signal_bus.subscriber import SignalSubscriber
+from asyncio_signal_bus.subscriber import BatchSignalSubscriber, SignalSubscriber
 from asyncio_signal_bus.types import R, S
 
 LOGGER = getLogger(__name__)
@@ -89,7 +89,52 @@ def subscriber(
         self._queues.get(topic_name).append(queue)
 
         def _wrapper(f: Callable[[S], Awaitable[R]]) -> SignalSubscriber[S, R]:
-            s = SignalSubscriber(error_handler(f), queue)
+            s = SignalSubscriber(
+                error_handler(f), queue, shutdown_timeout=shutdown_timeout
+            )
+            LOGGER.debug(f"Registering subscriber to topic {topic_name}")
+            self._subscribers.append(s)
+            return s
+
+        return _wrapper
+
+    def batch_subscriber(
+        self,
+        topic_name="default",
+        error_handler: Type[SubscriberErrorHandler] = SubscriberErrorHandler[S, R],
+        shutdown_timeout: Optional[SupportsFloat] = 120,
+        max_items: int = 10,
+        period_seconds: int = 10,
+    ):
+        """
+        A subscriber that consumes batches of events. The subscriber will wait no longer
+        than the period_seconds between aggregations. Batches will not exceed batch
+        seconds in size.
+        :param topic_name: The name of the topic used to link one or more subscribers
+            with one or more publishers.
+        :param error_handler: An error handler used to handle errors within the callable.
+            Error handling should usually terminate at the subscriber, with the
+            subscriber catching all exceptions. Any unhandled errors will block the
+            shutdown of the bus when the bus exits context or the stop method is used.
+        :param shutdown_timeout: If the subscriber takes longer than this time during
+            shutdown, then the task is killed and an error is raised. If you do not
+            want the task timeout to be limited, then set this value to None.
+        :param max_items: The maximum amount of time for the batch.
+        :param period_seconds: The maximum amount of timem to wait between batches.
+        :return: Wrapped callable
+        """
+        self._queues.setdefault(topic_name, [])
+        queue = Queue()
+        self._queues.get(topic_name).append(queue)
+
+        def _wrapper(f: Callable[[S], Awaitable[R]]) -> SignalSubscriber[S, R]:
+            s = BatchSignalSubscriber(
+                error_handler(f),
+                queue,
+                shutdown_timeout=shutdown_timeout,
+                max_items=max_items,
+                period_seconds=period_seconds,
+            )
             LOGGER.debug(f"Registering subscriber to topic {topic_name}")
             self._subscribers.append(s)
             return s

src/asyncio_signal_bus/subscriber.py

@@ -1,7 +1,8 @@
 import asyncio
 from asyncio import Queue, Task
 from logging import getLogger
-from typing import Awaitable, Callable, Generic, Optional, SupportsFloat
+from time import time
+from typing import Awaitable, Callable, Generic, Optional, SupportsFloat, SupportsInt
 
 from asyncio_signal_bus.exception import SignalBusShutdownError
 from asyncio_signal_bus.types import R, S
@@ -71,3 +72,42 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
     async def __call__(self, signal: S, *args, **kwargs) -> R:
         return await self._f(signal, *args, **kwargs)
+
+
+class BatchSignalSubscriber(SignalSubscriber):
+    def __init__(
+        self,
+        f: Callable[[S], Awaitable[R]],
+        queue: Queue,
+        shutdown_timeout: SupportsFloat = 120,
+        max_items: SupportsInt = 10,
+        period_seconds: SupportsFloat = 10,
+    ):
+        super().__init__(f, queue, shutdown_timeout)
+        self.max_items = int(max_items)
+        self.period_seconds = float(period_seconds)
+
+    async def _batch_task_wrapper(self, coroutine, n_items: int):
+        await coroutine
+        for i in range(n_items):
+            self._queue.task_done()
+
+    async def _listen(self):
+        LOGGER.debug("Started listening.")
+        batch = []
+        ts = time()
+        while True:
+            if len(batch) < self.max_items and not self._queue.empty():
+                signal = self._queue.get_nowait()
+                batch.append(signal)
+            if batch and (
+                len(batch) >= self.max_items or (time() - ts) > self.period_seconds
+            ):
+                asyncio.create_task(self._batch_task_wrapper(self(batch), len(batch)))
+                ts = time()
+                batch = []
+            # This deadlocks unless there is some sleep. It does not need to be long.
+            # My guess is that under the hood we just need to allow at least one clock
+            # cycle for something else to happen so that the event loop can perform
+            # other operations.
+            await asyncio.sleep(1e-10)

tests/test_subscriber.py

@@ -1,11 +1,11 @@
 import asyncio
 from asyncio import Queue
-from unittest.mock import Mock
+from unittest.mock import Mock, call
 
 import pytest
 
 from asyncio_signal_bus.exception import SignalBusShutdownError
-from asyncio_signal_bus.subscriber import SignalSubscriber
+from asyncio_signal_bus.subscriber import BatchSignalSubscriber, SignalSubscriber
 
 
 @pytest.mark.asyncio
@@ -36,3 +36,27 @@ async def foo_subscriber(signal: str):
     with pytest.raises(SignalBusShutdownError):
         async with signal_subscriber:
             await subscriber_queue.put("a")
+
+
+@pytest.mark.asyncio
+async def test_batch_subscriber():
+    subscriber_queue = Queue()
+
+    target_mock = Mock()
+
+    async def foo_subscriber(signal: str):
+        target_mock(signal)
+
+    signal_subscriber = BatchSignalSubscriber(
+        foo_subscriber,
+        subscriber_queue,
+        max_items=3,
+        period_seconds=0.1,
+        shutdown_timeout=0.3,
+    )
+    async with signal_subscriber:
+        for i in range(5):
+            await subscriber_queue.put(i)
+        await asyncio.sleep(0.2)
+        await subscriber_queue.put(6)
+    target_mock.assert_has_calls([call([0, 1, 2]), call([3, 4]), call([6])])