TRFC-2: Multiqueue support

[s3rius]

Abstract

Currently our brokers only leverage one queue and all tasks are sent in a single queue. The idea is to make it possible to assign different queues to tasks and being able to send tasks and read from multiple queues.

Solution

Currently there's a PR #418 that adds ability to add new methods to broker.task. We can make use of this feature in order to achieve the following API:

@broker.task.queue(...)
async def task(): ...

The question is how to specify queues, because queue API is different to all brokers.

One of possible solutions is to add function extra_queue method to broker, that returns a generic ExtraQueue object which can be used by the same broker to change task's destination.

Here's a usage example:

second_queue = broker.extra_queue(...)

# And then we can use these objects to assign tasks to these queues.

@broker.task.queue(second_queue)
async def task() :..

Questions

• Define ExtraQueue class. (or return strings)
• How to declare queues that brokers need to read from?
• Shall we split extra_queue to extra_read_queue and extra_write_queue

[vonsteer]

Before anything, thanks for continuing to improve this library! I'm particularly interested in this feature as I got a use case for it at work that is currently being handled (poorly) with celery that I would love to replace.

At least from my perspective working with SQS and on taskiq_aio_sqs:

• this kind of feature for SQS would require an extendable ExtraQueue class rather than strings as there's multiple configuration options for each queue (e.g. fifo, MessageVisibilityTimeout, MaxNumberOfMessages...).
• I personally wouldn't split between extra_read_queue/extra_write_queue, in the end the broker will be setting up a single connection to each queue so should be simple enough to just support extra_queue for both read/write internally. This is the way that celery handles it with broker_transport_options, defining the predefined queues once and then the task usage determines what is done with the queues.
• Maybe Queue object might be better naming, then you could also standardize the way the queue options are defined for each broker by each requiring at least one Queue object on broker __init__ that they can extend with more options depending on the broker.

Also, I was wondering with this feature, would we be able to define a single task that reads from a number of queues like so?

queues = [broker.extra_queue(...) i for range(1, 10)]

@broker.task.queues(*queues)
async def task() :..

it could simplify the definition for having a task consuming from a number of queues a lot.

[spikeninja]

@s3rius what about this?

IMHO answers to your questions:

1. 100% define Queue class
2. Let's read from all of them firstly, then this behavior can be extended easily
3. (the same answer as above)

Proposed syntax for the feature:

# redis_example.py
from taskiq import Queue
from taskiq_redis import RedisStreamBroker, StreamQueue

first = Queue(name="first")
second = Queue(name="second")
third = StreamQueue(name="critical", maxlen=1000)
broker = RedisStreamBroker(
    url="...",
    default_queue=first,
    queues={first, second, third},
)

@broker.task(queue_name=second.name)
async def funny_task(...): ...

funny_task.kiq(...)  # kicks to the queue which is provided in the decorator, namely "second"
funny_task.kiq(queue_name=third.name, ...)  # kicks to the provided queue, which is "third"

Proposed list of changes:

1. Define Queue class, for example:

@dataclass(slots=True, kw_only=True, frozen=True)
class Queue:
    name: str
    options: Mapping[str, Any] = field(default_factory=dict)
    def __str__(self) -> str: 
        return self.name

then we can easily declare specific queues in all adapters:

# taskiq_redis/queues.py
from dataclasses import dataclass
from taskiq.queues import Queue

@dataclass(slots=True, kw_only=True, frozen=True)
class RedisStreamQueue(Queue):
    maxlen: int | None = None
    approximate: bool = True

2. Add default_queue param to AsyncBroker to force all brokers to define it (not sure about that)
3. Add queue_name to TaskiqMessage and BrokerMessage entities
4. kick() method uses msg.queue_name to find out which queue to push the message
5. All code from listen() method is taken out to a private method __listen_single() (or give it any other name), then listen() method just runs concurrently (via asyncio.gather, etc) __listen_single() for each Queue in self.queues.

[s3rius]

I like the proposed changes. This will definitely work. Let's proceed with this approach. Additionally, regarding the question from @vonsteer, will it be able to read tasks from multiple queues?

Here's a small point to consider: currently, by default, we store all available tasks in the broker instance. If we implement this feature, it will be able to read tasks from all queues by default. Should we restrict it to execute tasks only if they come from a specific queue?

I would also suggest a slight modification to the API to avoid specifying common framework elements in the broker arguments. I propose adding queues to have an API as shown below.

Please also take a look at how to define a queue for a specific kiq. In this example, we use kicker in the same way as if we wanted to supply it with custom labels.

# redis_example.py
from taskiq import Queue
from taskiq_redis import RedisStreamBroker, StreamQueue

first = Queue(name="first")
second = Queue(name="second")
third = StreamQueue(name="critical", maxlen=1000)
broker = RedisStreamBroker(
    url="...",
).with_queues(first, second, third).with_default_queue(first)

@broker.task(queue_name=second.name)
async def funny_task(...): ...

funny_task.kiq(...)  # Kicks to the queue specified in the decorator, namely "second"
funny_task.kicker().with_queue(third.name).kiq()  # Kicks to the provided queue, which is "third"