Python asyncio.Queue在协程间通信的容量限制

asyncio.Queue 是 Python 异步编程中协程间通信的核心工具。它允许一个协程将数据放入队列，另一个协程从队列中取出数据，两者可以独立运行。asyncio.Queue 的一个关键特性是容量限制，即 maxsize 参数。这个限制直接关系到你的程序如何处理数据流，避免内存溢出，并控制生产者与消费者之间的速率匹配。

什么是 asyncio.Queue 和容量限制？

asyncio.Queue 是一个线程安全（在异步上下文中）的先进先出（FIFO）队列。它专为异步操作设计，提供了 put() 和 get() 方法，这些方法在队列满或空时会自动暂停协程，而不是阻塞整个事件循环。

容量限制 (maxsize) 定义了队列中可以容纳的最大元素数量。当你创建一个队列时，可以指定这个大小：

import asyncio

# 创建一个最多可容纳5个元素的队列
queue = asyncio.Queue(maxsize=5)

如果 maxsize 为0或负数，队列大小没有限制。

为什么需要容量限制？

想象一个场景：一个生产者协程以极快的速度生成数据，而消费者协程处理数据的速度很慢。如果没有容量限制，生产者产生的数据会不断堆积在内存中，最终导致内存耗尽（OOM - Out of Memory）。通过设置一个合理的 maxsize，你可以强制生产者在队列满时暂停，从而给消费者足够的时间去处理已存在的数据，达到一种速率平衡。

核心方法：put()、get()、put_nowait() 和 get_nowait()

为了有效利用容量限制，你需要了解 asyncio.Queue 的几个核心方法：

1. put(item): 将 item 放入队列。如果队列已满，当前协程会暂停，等待队列中出现空位。这是一个协程函数，需要 await。
2. get(): 从队列中取出一个元素。如果队列为空，当前协程会暂停，等待有新元素被放入。这也是一个协程函数，需要 await。
3. put_nowait(item): 尝试立即将 item 放入队列。如果队列已满，它会立即抛出 asyncio.QueueFull 异常，而不会等待。
4. get_nowait(): 尝试立即从队列中取出一个元素。如果队列为空，它会立即抛出 asyncio.QueueEmpty 异常，而不会等待。

在实际应用中，put_nowait() 和 get_nowait() 与容量限制配合使用，可以让你更精细地控制协程的行为，而不是简单地让它们无限期地等待。

实战演练：生产者-消费者模型

让我们通过一个经典的生产者-消费者模型来理解容量限制的实际应用。

目标

创建一个生产者协程，它不断生成数字。创建一个消费者协程，它从队列中取出数字并打印。我们将使用一个容量有限的队列来连接它们。

步骤 1: 导入 asyncio

import asyncio

步骤 2: 定义生产者协程

生产者协程将数字放入队列。我们使用 put_nowait() 并捕获 QueueFull 异常，以便在队列满时进行重试。

async def producer(queue: asyncio.Queue, max_items: int):
    """
    生产者协程，向队列中放入数据。
    :param queue: 异步队列
    :param max_items: 要生产的总项目数
    """
    for i in range(max_items):
        item = f"item_{i}"
        try:
            # 尝试立即放入数据
            queue.put_nowait(item)
            print(f"生产者: 放入 {item}")
        except asyncio.QueueFull:
            # 如果队列满了，等待一下再重试
            print("生产者: 队列已满，等待...")
            await asyncio.sleep(1)  # 等待1秒
            # 重试放入
            queue.put_nowait(item)
            print(f"生产者: 重试后放入 {item}")
        await asyncio.sleep(0.5)  # 模拟生产耗时
    print("生产者: 生产完成")

步骤 3: 定义消费者协程

消费者协程从队列中取出数据。我们使用 get_nowait() 并捕获 QueueEmpty 异常，以便在队列为空时进行重试。

async def consumer(queue: asyncio.Queue):
    """
    消费者协程，从队列中取出数据。
    :param queue: 异步队列
    """
    while True:
        try:
            # 尝试立即取出数据
            item = queue.get_nowait()
            print(f"消费者: 获取 {item}")
            queue.task_done()  # 标记任务完成
        except asyncio.QueueEmpty:
            # 如果队列为空，等待一下再重试
            print("消费者: 队列为空，等待...")
            await asyncio.sleep(1)  # 等待1秒
        await asyncio.sleep(1)  # 模拟处理耗时

步骤 4: 定义主函数并运行

在主函数中，我们创建队列、生产者和消费者，并启动它们。

async def main():
    # 创建一个容量为3的队列
    queue = asyncio.Queue(maxsize=3)

    # 创建生产者和消费者任务
    producer_task = asyncio.create_task(producer(queue, 10))
    consumer_task = asyncio.create_task(consumer(queue))

    # 等待生产者完成
    await producer_task

    # 等待队列中的所有项目都被处理
    await queue.join()

    # 取消消费者任务（因为它是一个无限循环）
    consumer_task.cancel()
    try:
        await consumer_task
    except asyncio.CancelledError:
        pass

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

运行结果分析

当你运行这段代码时，你会看到生产者在队列满（容量为3）时暂停，并打印“队列已满，等待...”。同时，消费者在处理数据，队列中的项目数会波动。这直观地展示了容量限制如何控制数据流，防止生产者过快。

如何优雅地停止协程？

上面的消费者协程是一个无限循环。在实际应用中，我们需要一种方法来通知消费者，当所有生产任务完成后，它可以安全地退出。最常用的方法是使用一个特殊的“哨兵值”（Sentinel Value），比如 None。

修改生产者协程

生产者在完成所有生产任务后，向队列中放入一个哨兵值。

async def producer(queue: asyncio.Queue, max_items: int):
    for i in range(max_items):
        item = f"item_{i}"
        try:
            queue.put_nowait(item)
            print(f"生产者: 放入 {item}")
        except asyncio.QueueFull:
            print("生产者: 队列已满，等待...")
            await asyncio.sleep(1)
            queue.put_nowait(item)
            print(f"生产者: 重试后放入 {item}")
        await asyncio.sleep(0.5)

    # 生产完成后，放入哨兵值以通知消费者
    print("生产者: 生产完成，放入哨兵值 None")
    queue.put_nowait(None) # 哨兵值

修改消费者协程

消费者在获取到哨兵值时，知道任务结束，并退出循环。

async def consumer(queue: asyncio.Queue):
    while True:
        item = await queue.get() # 使用 await get() 等待，直到有数据
        if item is None:
            # 收到哨兵值，退出循环
            print("消费者: 收到哨兵值，停止工作")
            break
        print(f"消费者: 获取 {item}")
        queue.task_done()
        await asyncio.sleep(1)

修改主函数

主函数不需要大的改动，queue.join() 会等待所有非哨兵值的项目都被处理完毕。

async def main():
    queue = asyncio.Queue(maxsize=3)
    producer_task = asyncio.create_task(producer(queue, 10))
    consumer_task = asyncio.create_task(consumer(queue))

    await producer_task
    await queue.join() # 等待所有非哨兵值的项目被处理
    consumer_task.cancel()
    try:
        await consumer_task
    except asyncio.CancelledError:
        pass

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

多消费者场景

如果有多个消费者，你需要放入与消费者数量相同的哨兵值，以确保每个消费者都能收到一个结束信号。

async def main():
    queue = asyncio.Queue(maxsize=3)
    num_consumers = 3
    producer_task = asyncio.create_task(producer(queue, 10))
    consumer_tasks = [asyncio.create_task(consumer(queue)) for _ in range(num_consumers)]

    await producer_task

    # 放入与消费者数量相同的哨兵值
    for _ in range(num_consumers):
        queue.put_nowait(None)

    await queue.join()
    for task in consumer_tasks:
        task.cancel()
    await asyncio.gather(*consumer_tasks, return_exceptions=True)

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

容量限制的权衡

选择合适的 maxsize 是一个需要权衡的过程。

最佳实践：从一个小容量开始，根据你的生产者和消费者的实际性能进行调优。监控内存使用情况和任务处理速率，找到最适合你应用的平衡点。