Aggregation of 2 RabbitMQ messages does not work properly (messages hanging unacked)

Question

I need to listen tasks on 2 queues, so I wrote the code below, but it has a problem. Currently it behaves like this: if the code started when 2 queues were full, it works great. But if queues were empty one of them was, the code reads messages, but does not proccess them (does not send ack, does not do the logic). But the messages became unacked, until I stop the code. I do not see any reason to be them unacked and unprocessed.

I can't understand what is wrong with the code? May be there is another way to aggregate 2 or more queues like this?

# task_processor.py

from aio_pika import IncomingMessage


class TaskProcessor:
    MAX_TASKS_PER_INSTANCE = 1

    def __init__(self):
        self._tasks = []

    def can_accept_new_task(self) -> bool:
        return len(self._tasks) < self.MAX_TASKS_PER_INSTANCE

    async def process(self, message: IncomingMessage):
        self._tasks.append(message)
        print(message.body)
        await message.ack()
        self._tasks.pop()

# main.py

import asyncio
from asyncio import QueueEmpty
from typing import Callable

import aio_pika
from aio_pika import RobustQueue
from dotenv import load_dotenv

load_dotenv()

from core.logger.logger import logger
from core.services.rabbitmq.task_processor.task_processor import TaskProcessor


async def get_single_task(queue: RobustQueue):
    while True:
        try:
            msg = await queue.get(timeout=3600)
            return msg
        except QueueEmpty:
            await asyncio.sleep(3)
        except asyncio.exceptions.TimeoutError:
            logger.warning('queue timeout error')
            pass
        except Exception as ex:
            logger.error(f"{queue} errored", exc_info=ex)


async def task_aggregator(queue1: RobustQueue, queue2: RobustQueue, should_take_new_task_cb: Callable):
    while True:
        if should_take_new_task_cb():
            queue2, queue1 = queue1, queue2
            gen1 = get_single_task(queue1)
            gen2 = get_single_task(queue2)

            done, _ = await asyncio.wait([gen1, gen2], return_when=asyncio.FIRST_COMPLETED)
            for item in done:
                result = item.result()
                yield result
        else:
            await asyncio.sleep(1)


async def tasks(queue1: RobustQueue, queue2: RobustQueue, should_take_new_task_cb: Callable):
    async for task in task_aggregator(queue1, queue2, should_take_new_task_cb):
        yield task


async def main():
    connection = await aio_pika.connect_robust(
        f"amqp://user:password@host:port/vhost?heartbeat={180}"
    )

    channel1 = connection.channel()
    channel2 = connection.channel()

    await channel1.initialize()
    await channel2.initialize()

    queue1 = await channel1.get_queue('queue1')
    queue2 = await channel2.get_queue('queue2')

    task_processor = TaskProcessor()
    task_generator = tasks(queue1, queue2, task_processor.can_accept_new_task)
    while True:
        if task_processor.can_accept_new_task():
            task = await anext(task_generator)
            await task_processor.process(task)
        else:
            await asyncio.sleep(1)


if __name__ == '__main__':
    asyncio.run(main())

Aggregation of 2 RabbitMQ messages does not work properly (messages hanging unacked)

0 Answers0