Changed listen signature for brokers. (#100)

Author: s3rius

docs/examples/extending/broker.py

@@ -25,12 +25,12 @@ async def shutdown(self) -> None:
         return await super().shutdown()
 
     async def kick(self, message: BrokerMessage) -> None:
-        # Send a message.
+        # Send a message.message.
         pass
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+    async def listen(self) -> AsyncGenerator[bytes, None]:
         while True:
             # Get new message.
-            new_message: BrokerMessage = ...  # type: ignore
+            new_message: bytes = ...  # type: ignore
             # Yield it!
             yield new_message

docs/extending-taskiq/broker.md

@@ -11,6 +11,20 @@ Here is a template for new brokers:
 
 @[code python](../examples/extending/broker.py)
 
+
+# About kick and listen
+
+The `kick` method takes a `BrokerMessage` as a parameter. The `BrokerMessage` class is a handy helper class for brokers. You can use information from the BrokerMessage to alter the delivery method.
+
+::: warning "cool warning!"
+
+As a broker developer, please send only raw bytes from the `message` field of a BrokerMessage if possible. Serializing it to the string may result in a problem if message bytes are not utf-8 compatible.
+
+:::
+
+
+The `listen` method should yield raw bytes that were sent over the network.
+
 ## Conventions
 
 For brokers, we have several conventions. It's good if your broker implements them.

taskiq/abc/broker.py

@@ -158,11 +158,14 @@ async def kick(
         Using this method tasks are sent to
         workers.
 
+        You don't need to send broker message. It's helper for brokers,
+        please send only bytes from message.message.
+
         :param message: name of a task.
         """
 
     @abstractmethod
-    def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+    def listen(self) -> AsyncGenerator[bytes, None]:
         """
         This function listens to new messages and yields them.
 

taskiq/abc/formatter.py

@@ -16,7 +16,7 @@ def dumps(self, message: TaskiqMessage) -> BrokerMessage:
         """
 
     @abstractmethod
-    def loads(self, message: BrokerMessage) -> TaskiqMessage:
+    def loads(self, message: bytes) -> TaskiqMessage:
         """
         Parses broker message to TaskiqMessage.
 

taskiq/brokers/inmemory_broker.py

@@ -141,11 +141,11 @@ async def kick(self, message: BrokerMessage) -> None:
                 target_task.original_func,
             )
 
-        task = asyncio.create_task(self.receiver.callback(message=message))
+        task = asyncio.create_task(self.receiver.callback(message=message.message))
         self._running_tasks.add(task)
         task.add_done_callback(self._running_tasks.discard)
 
-    def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+    def listen(self) -> AsyncGenerator[bytes, None]:
         """
         Inmemory broker cannot listen.
 

taskiq/brokers/shared_broker.py

@@ -59,7 +59,7 @@ async def kick(self, message: BrokerMessage) -> None:
             "without setting the default_broker.",
         )
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:  # type: ignore
+    async def listen(self) -> AsyncGenerator[bytes, None]:  # type: ignore
         """
         Shared broker cannot listen to tasks.
 

taskiq/brokers/zmq_broker.py

@@ -1,3 +1,4 @@
+import math
 from logging import getLogger
 from typing import AsyncGenerator, Callable, Optional, TypeVar
 
@@ -58,21 +59,23 @@ async def kick(self, message: BrokerMessage) -> None:
 
         :param message: message to publish.
         """
+        part_len = 100
+        parts = [
+            message.message[
+                idx * part_len : min(idx * part_len + part_len, len(message.message))
+            ]
+            for idx in range(math.ceil(len(message.message) / part_len))
+        ]
         with self.socket.connect(self.sub_host) as sock:
-            await sock.send_string(message.json())
+            await sock.send_multipart(parts)
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+    async def listen(self) -> AsyncGenerator[bytes, None]:
         """
         Start accepting new messages.
 
         :yields: incoming messages.
         """
         with self.socket.connect(self.sub_host) as sock:
-            while True:
-                received_str = await sock.recv_string()
-                try:
-                    broker_msg = BrokerMessage.parse_raw(received_str)
-                except ValueError:
-                    logger.warning("Cannot parse received message %s", received_str)
-                    continue
-                yield broker_msg
+            while True:  # noqa: WPS457
+                data = await sock.recv_multipart()
+                yield b"".join(data)

taskiq/formatters/json_formatter.py

@@ -19,11 +19,11 @@ def dumps(self, message: TaskiqMessage) -> BrokerMessage:
             labels=message.labels,
         )
 
-    def loads(self, message: BrokerMessage) -> TaskiqMessage:
+    def loads(self, message: bytes) -> TaskiqMessage:
         """
         Loads json from message.
 
         :param message: broker's message.
         :return: parsed taskiq message.
         """
-        return TaskiqMessage.parse_raw(message.message)
+        return TaskiqMessage.parse_raw(message)

taskiq/receiver/receiver.py

@@ -10,7 +10,7 @@
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.context import Context
-from taskiq.message import BrokerMessage, TaskiqMessage
+from taskiq.message import TaskiqMessage
 from taskiq.receiver.params_parser import parse_params
 from taskiq.result import TaskiqResult
 from taskiq.state import TaskiqState
@@ -64,7 +64,7 @@ def __init__(
 
     async def callback(  # noqa: C901, WPS213
         self,
-        message: BrokerMessage,
+        message: bytes,
         raise_err: bool = False,
     ) -> None:
         """
@@ -79,17 +79,6 @@ async def callback(  # noqa: C901, WPS213
         :param raise_err: raise an error if cannot save result in
             result_backend.
         """
-        logger.debug(f"Received message: {message}")
-        if message.task_name not in self.broker.available_tasks:
-            logger.warning(
-                'task "%s" is not found. Maybe you forgot to import it?',
-                message.task_name,
-            )
-            return
-        logger.debug(
-            "Function for task %s is resolved. Executing...",
-            message.task_name,
-        )
         try:
             taskiq_msg = self.broker.formatter.loads(message=message)
         except Exception as exc:
@@ -100,6 +89,17 @@ async def callback(  # noqa: C901, WPS213
                 exc_info=True,
             )
             return
+        logger.debug(f"Received message: {taskiq_msg}")
+        if taskiq_msg.task_name not in self.broker.available_tasks:
+            logger.warning(
+                'task "%s" is not found. Maybe you forgot to import it?',
+                taskiq_msg.task_name,
+            )
+            return
+        logger.debug(
+            "Function for task %s is resolved. Executing...",
+            taskiq_msg.task_name,
+        )
         for middleware in self.broker.middlewares:
             if middleware.__class__.pre_execute != TaskiqMiddleware.pre_execute:
                 taskiq_msg = await maybe_awaitable(
@@ -114,14 +114,14 @@ async def callback(  # noqa: C901, WPS213
             taskiq_msg.task_id,
         )
         result = await self.run_task(
-            target=self.broker.available_tasks[message.task_name].original_func,
+            target=self.broker.available_tasks[taskiq_msg.task_name].original_func,
             message=taskiq_msg,
         )
         for middleware in self.broker.middlewares:
             if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
                 await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
         try:
-            await self.broker.result_backend.set_result(message.task_id, result)
+            await self.broker.result_backend.set_result(taskiq_msg.task_id, result)
             for middleware in self.broker.middlewares:
                 if middleware.__class__.post_save != TaskiqMiddleware.post_save:
                     await maybe_awaitable(middleware.post_save(taskiq_msg, result))

tests/cli/worker/test_receiver.py

@@ -9,7 +9,7 @@
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.brokers.inmemory_broker import InMemoryBroker
-from taskiq.message import BrokerMessage, TaskiqMessage
+from taskiq.message import TaskiqMessage
 from taskiq.receiver import Receiver
 from taskiq.result import TaskiqResult
 
@@ -28,9 +28,9 @@ def __init__(
         )
         self.to_send: "List[TaskiqMessage]" = []
 
-    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
+    async def listen(self) -> AsyncGenerator[bytes, None]:
         for message in self.to_send:
-            yield self.formatter.dumps(message)
+            yield self.formatter.dumps(message).message
 
 
 def get_receiver(
@@ -186,7 +186,7 @@ async def my_task() -> int:
         ),
     )
 
-    await receiver.callback(broker_message)
+    await receiver.callback(broker_message.message)
     assert called_times == 1
 
 
@@ -196,12 +196,7 @@ async def test_callback_wrong_format() -> None:
     receiver = get_receiver()
 
     await receiver.callback(
-        BrokerMessage(
-            task_id="",
-            task_name="my_task.task_name",
-            message='{"aaaa": "bbb"}',
-            labels={},
-        ),
+        b"{some wrong bytes}",
     )
 
 
@@ -221,7 +216,7 @@ async def test_callback_unknown_task() -> None:
         ),
     )
 
-    await receiver.callback(broker_message)
+    await receiver.callback(broker_message.message)
 
 
 @pytest.mark.anyio