在 Asyncio 中兼顾生产者和消费者 Python

Question

I have a producer and a consumer async function. The producer runs first and then three consumers consume the task concurrently. I want to continue this cadence indefinitely.

• What is the best way to stop the producer and send a signal to the consumer to start running?
• What is the best way to stop the consumer and send a signal to the producer to start running?

My current setup looks like this:

import asyncio
import time


async def producer(event):
    n = 0
    while True:
        print("Running producer...")
        await asyncio.sleep(0.5)
        n += 1
        if n == 2:
            event.set()
            break


async def consumer(event):
    await event.wait()
    print("Running consumer...")
    await asyncio.sleep(0.5)


async def main():
    event = asyncio.Event()

    tasks = [asyncio.create_task(producer(event))] + [
        asyncio.create_task(consumer(event)) for _ in range(3)
    ]

    await asyncio.gather(*tasks)


while True:
    asyncio.run(main())
    print("\nSleeping for 1 sec...\n")
    time.sleep(1)

This produces the following output:

Running producer...
Running producer...
Running consumer...
Running consumer...
Running consumer...

Sleeping for 1 sec...

Running producer...
Running producer...
Running consumer...
Running consumer...
Running consumer...

The snippet above will run the producer and the two consumers indefinitely. This works as expected:

• The producer and the two consumers run concurrently
• The while loop surrounding asyncio.run is running the system indefinitely

However, I was wondering if there is a better synchronization technique to achieve this long-running periodicity?

Answer 1

An alternative to using an event to wake up your consumers is to synchronize the producer and consumers using a queue. The code would look something like

import asyncio


async def producer(queue: asyncio.Queue):
    while True:
        print("Running producer...")
        message = await fetch_message_from_sqs()
        if message:
            await queue.put(message)
        await asyncio.sleep(0.5)


async def consumer(queue: asyncio.Queue):
    while True:
        print("Running consumer...")
        if queue.empty():
            await asyncio.sleep(1)
            continue
        message = await queue.get()
        print(message)
        await acknowledge_message(message)


async def main():
    # You can set a max size if you want to prevent pull too many messages from SQS.
    queue = asyncio.Queue()

    tasks = asyncio.gather(producer(queue), *[consumer(queue) for _ in range(3)])


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