Merge pull request #12 from taskiq-python/improve_jetstream_broker

Author: s3rius

.github/workflows/test.yaml

@@ -1,4 +1,4 @@
-name: Testing taskiq-redis
+name: Testing taskiq-nats
 
 on: pull_request
 

README.md

@@ -15,6 +15,7 @@ pip install taskiq taskiq-nats
 
 Here's a minimal setup example with a broker and one task.
 
+### Default NATS broker.
 ```python
 import asyncio
 from taskiq_nats import NatsBroker, JetStreamBroker
@@ -27,15 +28,47 @@ broker = NatsBroker(
     queue="random_queue_name",
 )
 
-# Or alternatively you can use a JetStream broker:
-broker = JetStreamBroker(
-    [
+
+@broker.task
+async def my_lovely_task():
+    print("I love taskiq")
+
+
+async def main():
+    await broker.startup()
+
+    await my_lovely_task.kiq()
+
+    await broker.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+
+```
+### NATS broker based on JetStream
+```python
+import asyncio
+from taskiq_nats import (
+    PushBasedJetStreamBroker,
+    PullBasedJetStreamBroker
+)
+
+broker = PushBasedJetStreamBroker(
+    servers=[
         "nats://nats1:4222",
         "nats://nats2:4222",
     ],
-    queue="random_queue_name",
-    subject="my-subj",
-    stream_name="my-stream"
+    queue="awesome_queue_name",
+)
+
+# Or you can use pull based variant
+broker = PullBasedJetStreamBroker(
+    servers=[
+        "nats://nats1:4222",
+        "nats://nats2:4222",
+    ],
+    durable="awesome_durable_consumer_name",
 )
 
 
@@ -54,7 +87,6 @@ async def main():
 
 if __name__ == "__main__":
     asyncio.run(main())
-
 ```
 
 ## NatsBroker configuration
@@ -68,3 +100,22 @@ Here's the constructor parameters:
 * `result_backend` - custom result backend.
 * `task_id_generator` - custom function to generate task ids.
 * Every other keyword argument will be sent to `nats.connect` function.
+
+## JetStreamBroker configuration
+### Common
+* `servers` - a single string or a list of strings with nats nodes addresses.
+* `subject` - name of the subect that will be used to exchange tasks betwee workers and clients.
+* `stream_name` - name of the stream where subjects will be located.
+* `queue` - a single string or a list of strings with nats nodes addresses.
+* `result_backend` - custom result backend.
+* `task_id_generator` - custom function to generate task ids.
+* `stream_config` - a config for stream.
+* `consumer_config` - a config for consumer.
+
+### PushBasedJetStreamBroker
+* `queue` - name of the queue. It's used to share messages between different consumers.
+
+### PullBasedJetStreamBroker
+* `durable` - durable name of the consumer. It's used to share messages between different consumers.
+* `pull_consume_batch` - maximum number of message that can be fetched each time.
+* `pull_consume_timeout` - timeout for messages fetch. If there is no messages, we start fetching messages again.

poetry.lock

@@ -5,6 +5,14 @@
 uses NATS as a message queue.
 """
 
-from taskiq_nats.broker import JetStreamBroker, NatsBroker
+from taskiq_nats.broker import (
+    NatsBroker,
+    PullBasedJetStreamBroker,
+    PushBasedJetStreamBroker,
+)
 
-__all__ = ["NatsBroker", "JetStreamBroker"]
+__all__ = [
+    "NatsBroker",
+    "PushBasedJetStreamBroker",
+    "PullBasedJetStreamBroker",
+]

taskiq_nats/__init__.py

@@ -1,12 +1,20 @@
+import typing
+from abc import ABC, abstractmethod
 from logging import getLogger
-from typing import Any, AsyncGenerator, Callable, List, Optional, TypeVar, Union
 
 from nats.aio.client import Client
+from nats.aio.msg import Msg as NatsMessage
+from nats.errors import TimeoutError as NatsTimeoutError
 from nats.js import JetStreamContext
-from nats.js.api import StreamConfig
+from nats.js.api import ConsumerConfig, StreamConfig
 from taskiq import AckableMessage, AsyncBroker, AsyncResultBackend, BrokerMessage
 
-_T = TypeVar("_T")  # noqa: WPS111 (Too short)
+_T = typing.TypeVar("_T")  # noqa: WPS111 (Too short)
+
+
+JetStreamConsumerType = typing.TypeVar(
+    "JetStreamConsumerType",
+)
 
 
 logger = getLogger("taskiq_nats")
@@ -29,12 +37,12 @@ class NatsBroker(AsyncBroker):
 
     def __init__(  # noqa: WPS211 (too many args)
         self,
-        servers: Union[str, List[str]],
+        servers: typing.Union[str, typing.List[str]],
         subject: str = "taskiq_tasks",
-        queue: Optional[str] = None,
-        result_backend: "Optional[AsyncResultBackend[_T]]" = None,
-        task_id_generator: Optional[Callable[[], str]] = None,
-        **connection_kwargs: Any,
+        queue: typing.Optional[str] = None,
+        result_backend: "typing.Optional[AsyncResultBackend[_T]]" = None,
+        task_id_generator: typing.Optional[typing.Callable[[], str]] = None,
+        **connection_kwargs: typing.Any,
     ) -> None:
         super().__init__(result_backend, task_id_generator)
         self.servers = servers
@@ -64,7 +72,7 @@ async def kick(self, message: BrokerMessage) -> None:
             headers=message.labels,
         )
 
-    async def listen(self) -> AsyncGenerator[bytes, None]:
+    async def listen(self) -> typing.AsyncGenerator[bytes, None]:
         """
         Start listen to new messages.
 
@@ -80,37 +88,55 @@ async def shutdown(self) -> None:
         await super().shutdown()
 
 
-class JetStreamBroker(AsyncBroker):  # noqa: WPS230
-    """
-    JetStream broker for taskiq.
+class BaseJetStreamBroker(  # noqa: WPS230 (too many attrs)
+    AsyncBroker,
+    ABC,
+    typing.Generic[JetStreamConsumerType],
+):
+    """Base JetStream broker for taskiq.
 
-    This broker creates a JetStream context
-    and uses it to send and receive messages.
+    It has two subclasses - PullBasedJetStreamBroker
+    and PushBasedJetStreamBroker.
+
+    These brokers create a JetStream context
+    and use it to send and receive messages.
 
     This is useful for systems where you need to
     be sure that messages are delivered to the workers.
     """
 
     def __init__(  # noqa: WPS211 (too many args)
         self,
-        servers: Union[str, List[str]],
-        subject: str = "tasiq_tasks",
-        stream_name: str = "taskiq_jstream",
-        queue: Optional[str] = None,
-        result_backend: "Optional[AsyncResultBackend[_T]]" = None,
-        task_id_generator: Optional[Callable[[], str]] = None,
-        stream_config: Optional[StreamConfig] = None,
-        **connection_kwargs: Any,
+        servers: typing.Union[str, typing.List[str]],
+        subject: str = "taskiq_tasks",
+        stream_name: str = "taskiq_jetstream",
+        queue: typing.Optional[str] = None,
+        durable: str = "taskiq_durable",
+        stream_config: typing.Optional[StreamConfig] = None,
+        consumer_config: typing.Optional[ConsumerConfig] = None,
+        pull_consume_batch: int = 1,
+        pull_consume_timeout: typing.Optional[float] = None,
+        **connection_kwargs: typing.Any,
     ) -> None:
-        super().__init__(result_backend, task_id_generator)
-        self.servers = servers
-        self.client: Client = Client()
-        self.connection_kwargs = connection_kwargs
-        self.queue = queue
-        self.subject = subject
-        self.stream_name = stream_name
+        super().__init__()
+        self.servers: typing.Final = servers
+        self.client: typing.Final = Client()
+        self.connection_kwargs: typing.Final = connection_kwargs
+        self.subject: typing.Final = subject
+        self.stream_name: typing.Final = stream_name
         self.js: JetStreamContext
         self.stream_config = stream_config or StreamConfig()
+        self.consumer_config = consumer_config
+
+        # Only for push based consumer
+        self.queue: typing.Final = queue
+        self.default_consumer_name: typing.Final = "taskiq_consumer"
+        # Only for pull based consumer
+        self.durable: typing.Final = durable
+        self.pull_consume_batch: typing.Final = pull_consume_batch
+        self.pull_consume_timeout: typing.Final = pull_consume_timeout
+
+        self.consumer: JetStreamConsumerType
 
     async def startup(self) -> None:
         """
@@ -127,6 +153,12 @@ async def startup(self) -> None:
         if not self.stream_config.subjects:
             self.stream_config.subjects = [self.subject]
         await self.js.add_stream(config=self.stream_config)
+        await self._startup_consumer()
+
+    async def shutdown(self) -> None:
+        """Close connections to NATS."""
+        await self.client.close()
+        await super().shutdown()
 
     async def kick(self, message: BrokerMessage) -> None:
         """
@@ -140,20 +172,89 @@ async def kick(self, message: BrokerMessage) -> None:
             headers=message.labels,
         )
 
-    async def listen(self) -> AsyncGenerator[AckableMessage, None]:
+    @abstractmethod
+    async def _startup_consumer(self) -> None:
+        """Create consumer."""
+
+
+class PushBasedJetStreamBroker(
+    BaseJetStreamBroker[JetStreamContext.PushSubscription],
+):
+    """JetStream broker for push based message consumption.
+
+    It's named `push` based because nats server push messages to
+    the consumer, not consumer requests them.
+    """
+
+    async def listen(self) -> typing.AsyncGenerator[AckableMessage, None]:
         """
         Start listen to new messages.
 
         :yield: incoming messages.
         """
-        subscribe = await self.js.subscribe(self.subject, queue=self.queue or "")
-        async for message in subscribe.messages:
+        async for message in self.consumer.messages:
             yield AckableMessage(
                 data=message.data,
                 ack=message.ack,
             )
 
-    async def shutdown(self) -> None:
-        """Close connections to NATS."""
-        await self.client.close()
-        await super().shutdown()
+    async def _startup_consumer(self) -> None:
+        if not self.consumer_config:
+            self.consumer_config = ConsumerConfig(
+                name=self.default_consumer_name,
+                durable_name=self.default_consumer_name,
+            )
+
+        self.consumer = await self.js.subscribe(
+            subject=self.subject,
+            queue=self.queue or "",
+            config=self.consumer_config,
+        )
+
+
+class PullBasedJetStreamBroker(
+    BaseJetStreamBroker[JetStreamContext.PullSubscription],
+):
+    """JetStream broker for pull based message consumption.
+
+    It's named `pull` based because consumer requests messages,
+    not NATS server sends them.
+    """
+
+    async def listen(self) -> typing.AsyncGenerator[AckableMessage, None]:
+        """
+        Start listen to new messages.
+
+        :yield: incoming messages.
+        """
+        while True:  # noqa: WPS327
+            try:
+                nats_messages: typing.List[NatsMessage] = await self.consumer.fetch(
+                    batch=self.pull_consume_batch,
+                    timeout=self.pull_consume_timeout,
+                )
+                for nats_message in nats_messages:
+                    yield AckableMessage(
+                        data=nats_message.data,
+                        ack=nats_message.ack,
+                    )
+            except NatsTimeoutError:
+                continue
+
+    async def _startup_consumer(self) -> None:
+        if not self.consumer_config:
+            self.consumer_config = ConsumerConfig(
+                durable_name=self.durable,
+            )
+
+        # We must use this method to create pull based consumer
+        # because consumer config won't change without it.
+        await self.js.add_consumer(
+            stream=self.stream_config.name or self.stream_name,
+            config=self.consumer_config,
+        )
+        self.consumer = await self.js.pull_subscribe(
+            subject=self.subject,
+            durable=self.durable,
+            config=self.consumer_config,
+        )

taskiq_nats/broker.py

@@ -1,8 +1,10 @@
 import os
 import uuid
-from typing import List
+from typing import AsyncGenerator, Final, List
 
 import pytest
+from nats import NATS
+from nats.js import JetStreamContext
 
 
 @pytest.fixture(scope="session")
@@ -38,3 +40,20 @@ def nats_urls() -> List[str]:
     """
     urls = os.environ.get("NATS_URLS") or "nats://localhost:4222"
     return urls.split(",")
+
+
+@pytest.fixture()
+async def nats_jetstream(
+    nats_urls: List[str],  # noqa: WPS442
+) -> AsyncGenerator[JetStreamContext, None]:
+    """Create and yield nats client and jetstream instances.
+
+    :param nats_urls: urls to nats.
+
+    :yields: NATS JetStream.
+    """
+    nats: Final = NATS()
+    await nats.connect(servers=nats_urls)
+    jetstream: Final = nats.jetstream()
+    yield jetstream
+    await nats.close()