RabbitMQ多线程处理消息：消息数不足线程数的问题修复

问题描述 #

使用多线程处理RabbitMQ消息队列时，当消息数量小于设定的最大线程数（最多5个线程），会出现消息残留：队列中有13个URL时，程序结束后剩余3个未处理；12个URL时剩余2个未处理。需求为：消息数≥5时使用5个工作线程，消息数<5时线程数等于消息数，确保队列无消息残留。

原代码 #

import asyncio
import aio_pika
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging
import traceback

logger = logging.getLogger(__name__)
messagesList = []
linkList = []
channel = None
queue = None
resultQuequ = None
resultExchange = None

def pageCollect(messageItem):
    # 假设这是你的页面收集逻辑，返回处理结果字典或None
    return {"url": "example.com"}

async def processTask(message: aio_pika.abc.AbstractIncomingMessage):
    try:        
        messagesList.append(message)
        # when message number =5
        if len(messagesList) == 5:
            messages_to_process = messagesList[:]  
            
            max_workers = 5 if len(messages_to_process) >= 5 else len(messages_to_process)
            with ThreadPoolExecutor(max_workers=max_workers) as pool:
                futures = {pool.submit(pageCollect, messageItem): messageItem for messageItem in
                           messages_to_process}
                for future in as_completed(futures):
                    messageItem = futures[future]
                    try:
                        resultDict = future.result()
                        if resultDict:
                            await messageItem.ack()
                            
                            logger.info(f'{messageItem.delivery_tag}: {resultDict}')
                        else:
                            await messageItem.reject(requeue=True)
                    except Exception as e:
                        logger.error(f'{messageItem.delivery_tag}: {e}')
                        await messageItem.reject(requeue=True)
            linkList.clear()
            messagesList.clear()

        # when message number < 5
        remaining_messages = queue.declaration_result.message_count
        if remaining_messages > 0:
            messages_to_process = messagesList[:]
            max_workers = len(messages_to_process)
            with ThreadPoolExecutor(max_workers=max_workers) as pool:
                futures = {pool.submit(pageCollect, messageItem): messageItem for messageItem in messagesList}
                for future in as_completed(futures):
                    messageItem = messages_to_process.pop(0)
                    try:
                        resultDict = future.result()
                        if resultDict:
                            await messageItem.ack()
                            logger.info(f'{messageItem.delivery_tag}: {resultDict}')
                        else:
                            await messageItem.reject(requeue=True)
                    except Exception as e:
                        logger.error(f'{messageItem.delivery_tag}: {e}')
                        await messageItem.reject(requeue=True)
            messagesList.clear()

        await asyncio.sleep(0.1)

    except Exception as e:
        print(e)


async def main(loop):
    try:
        # connect
        connection = await aio_pika.connect_robust(host='XX.XX.X.XX', port=5672, login='admin', password='admin',
                                                   virtualhost='my_vhost', loop=loop)
        
        global channel
        channel = await connection.channel()
        # Will take no more than 10 messages in advance
        await channel.set_qos() #prefetch_count=5
        crawler_exchange = await channel.declare_exchange(name='crawler_exchange', type='fanout')
        
        queueName = "myqueue"
        global queue
        queue = await channel.declare_queue(queueName, durable=True)
        await queue.bind(crawler_exchange, routing_key="myqueue")

        
        rstqueueName = "Result"
        global resultQuequ
        resultQuequ = await channel.declare_queue(rstqueueName, durable=True)
        
        global resultExchange
        resultExchange = await channel.declare_exchange(name='resultExchange', type='direct')
        
        await resultQuequ.bind(resultExchange, routing_key="allResult")

        # get message
        await queue.consume(processTask)
        logger.info(f"Waiting for messages at {queue.name}. To exit press CTRL+C")
        return connection

    except Exception as e:
        logger.error(f"failed： {e}")
        logger.error(traceback.format_exc())


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    connection = loop.run_until_complete(main(loop))

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        logger.info("Received exit signal")
    finally:
        loop.run_until_complete(connection.close())
        loop.close()

问题分析与修复方案 #

核心问题点 #

1. QoS配置未生效：原代码中await channel.set_qos()未指定prefetch_count，无法控制预取消息数量，导致消息分发逻辑混乱。
2. 剩余消息判断不准确：queue.declaration_result.message_count是队列声明时的初始数量，并非实时剩余消息数，无法作为触发处理的依据。
3. 重复处理逻辑冲突：当消息数不足5时，每条消息都会触发一次处理逻辑，导致缓存消息被重复提交到线程池，同时遗漏了队列耗尽后的剩余消息。
4. 无队列耗尽后的收尾处理：当队列中没有新消息时，缓存中剩余的不足5条的消息无法被触发处理，导致残留。

修复后的代码 #

import asyncio
import aio_pika
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging
import traceback

logger = logging.getLogger(__name__)
messagesList = []
linkList = []
channel = None
queue = None
resultQuequ = None
resultExchange = None
is_processing = False  # 防止并发处理冲突

def pageCollect(messageItem):
    # 替换为实际的页面收集逻辑
    try:
        # 模拟处理时间
        asyncio.sleep(0.5)
        return {"url": messageItem.body.decode(), "status": "success"}
    except Exception as e:
        logger.error(f"处理失败: {e}")
        return None

async def process_batch():
    """统一批量处理消息的逻辑"""
    global is_processing
    if is_processing or not messagesList:
        return
    is_processing = True
    
    messages_to_process = messagesList[:]
    messagesList.clear()
    
    max_workers = min(5, len(messages_to_process))
    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {pool.submit(pageCollect, msg): msg for msg in messages_to_process}
        for future in as_completed(futures):
            msg = futures[future]
            try:
                resultDict = future.result()
                if resultDict:
                    await msg.ack()
                    logger.info(f'{msg.delivery_tag}: {resultDict}')
                else:
                    await msg.reject(requeue=True)
            except Exception as e:
                logger.error(f'{msg.delivery_tag}: {e}')
                await msg.reject(requeue=True)
    is_processing = False

async def check_remaining_messages():
    """定时检查队列是否为空，处理剩余消息"""
    while True:
        await asyncio.sleep(1)
        # 获取实时队列消息数
        queue_status = await queue.get_declare_ok()
        remaining = queue_status.message_count
        if remaining == 0 and messagesList:
            await process_batch()

async def processTask(message: aio_pika.abc.AbstractIncomingMessage):
    try:        
        messagesList.append(message)
        # 缓存满5条时触发批量处理
        if len(messagesList) >= 5:
            await process_batch()
        await asyncio.sleep(0.1)
    except Exception as e:
        logger.error(f"处理消息时出错: {e}")
        await message.reject(requeue=True)

async def main(loop):
    try:
        connection = await aio_pika.connect_robust(
            host='XX.XX.X.XX', 
            port=5672, 
            login='admin', 
            password='admin',
            virtualhost='my_vhost', 
            loop=loop
        )
        
        global channel
        channel = await connection.channel()
        # 设置预取数为5，配合批量处理
        await channel.set_qos(prefetch_count=5)
        
        crawler_exchange = await channel.declare_exchange(name='crawler_exchange', type='fanout')
        queueName = "myqueue"
        global queue
        queue = await channel.declare_queue(queueName, durable=True)
        await queue.bind(crawler_exchange, routing_key="myqueue")

        rstqueueName = "Result"
        global resultQuequ
        resultQuequ = await channel.declare_queue(rstqueueName, durable=True)
        global resultExchange
        resultExchange = await channel.declare_exchange(name='resultExchange', type='direct')
        await resultQuequ.bind(resultExchange, routing_key="allResult")

        # 启动定时检查任务
        loop.create_task(check_remaining_messages())
        
        await queue.consume(processTask)
        logger.info(f"Waiting for messages at {queue.name}. To exit press CTRL+C")
        return connection

    except Exception as e:
        logger.error(f"初始化失败： {e}")
        logger.error(traceback.format_exc())

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    loop = asyncio.get_event_loop()
    connection = loop.run_until_complete(main(loop))

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        logger.info("Received exit signal")
    finally:
        # 退出前处理剩余消息
        loop.run_until_complete(process_batch())
        loop.run_until_complete(connection.close())
        loop.close()

关键修改说明 #

1. 统一批量处理逻辑：将分散的处理逻辑合并到process_batch函数，避免重复代码，同时用is_processing标志防止并发处理冲突。
2. 正确配置QoS：设置prefetch_count=5，确保通道最多预取5条消息，配合批量处理逻辑，避免消息堆积在客户端。
3. 实时获取队列剩余消息：使用await queue.get_declare_ok()获取实时队列状态，替代原有的过时数据。
4. 定时检查收尾：新增check_remaining_messages定时任务，当队列无消息且缓存有剩余时，触发批量处理，确保无残留。
5. 退出前收尾：在程序退出时，主动调用process_batch处理剩余缓存消息，避免退出时残留未处理消息。

内容的提问来源于stack exchange，提问作者RR l