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"

[spikeninja]

Agreed, the syntax above is much more convenient)

[spikeninja]

So, to sum up:

Targeted Syntax:

# 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"

Steps:

1. taskiq update (0.12.0)

   1. Define Queue class, for example:

   @dataclass(slots=True, kw_only=True, frozen=True)
   class Queue:
      name: str
      # I'm not sure about extra options, but just in case I will leave it here
      options: Mapping[str, Any] = field(default_factory=dict)
      def __str__(self) -> str: 
         return self.name

   2. Add with_queue(queue: Queue) to AsyncKicker
   3. Add default_queue, with_queues(*queues: Queue) and with_default_queue to AsyncBroker
   4. TaskiqMessage and BrokerMessage should encapsulate the corresponding queue names. Add queue_name to the both.

2. adapters update (taskiq>=0.12.0)

   1. kick() methods use msg.queue_name to find out which queue to push the message
   2. 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.

3. tests:

   1. @broker.task(queue_name="X") sends to "X" queue
   2. heart.kicker().with_queue("Y").kiq() sends to "Y" queue
   3. broker.with_queues(first, second) implicit default set (SHOULD BE DISCUSSED)
   4. queue does not exist
   5. ...

[s3rius]

I like the design. But have several small things to mention.

1. In case if no queues will be specified, shall we just declare our own?
2. In case if queues were specified, but no default_queue was selected, which one should be used by default. We can declare default queue anyway and use it as default unless with_default is specified.
3. In section 2.3 I would suggest rewriting it using anyio.group. Because it will be easier to maintain. Since it will cancel other tasks if any of them will ever fail.
4. What do you think about renaming task argument from queue_name to queue and accept Union[Queue, str]?

Other questions seems to be solved already. Looks solid. Thanks.

[spikeninja]

IMHO

1. I think we should (Queue("taskiq"))
2. "We can declare..." - Agreed 🤝
3. Agreed 🤝
4. Yes, I actually wanted to propose Union[Queue, str] initially but then for some reason changed my mind :) Let's go with Union[Queue, str]

@vonsteer your thoughts?

[vonsteer]

I think this design should be extendable enough to work on taskiq-aio-sqs and other brokers like a charm.

My only lingering doubt is queue based priority listening, do we want to consider this use case for brokers without native server support for message based priorities? Something like this, but purely for listening part of the broker with this kind of syntax:

high_priority = Queue(name="critical", priority=5)
medium_priority = Queue(name="normal", priority=2) 
low_priority = Queue(name="low", priority=1)

broker = RedisStreamBroker(url="redis://localhost:6379")
broker.with_queues(high_priority, medium_priority, low_priority)
    
await broker.listen(with_priority=True)

IMO this may be a bit out of scope for this RFC, but I do see it as a possible feature I'd like to implement in the future for taskiq-aio-sqs (and celery does not support it for the SQS broker).