當消費者在 asyncio.Queue 生產-消費流程中引發異常時，通知生產者停止生產

Question

基於asyncio.Queue的模擬生產-消費流程：

import asyncio


async def produce(q: asyncio.Queue, task):
    asyncio.create_task(q.put(task))
    print(f'Produced {task}')


async def consume(q: asyncio.Queue):
    while True:
        task = await q.get()
        if task > 2:
            print(f'Cannot consume {task}')
            raise ValueError(f'{task} too big')
        print(f'Consumed {task}')
        q.task_done()


async def main():
    queue = asyncio.Queue()
    consumers = [asyncio.create_task(consume(queue)) for _ in range(2)]
    for i in range(10):
        await asyncio.create_task(produce(queue, i))
    await asyncio.wait([queue.join(), *consumers],
                       return_when=asyncio.FIRST_COMPLETED)


asyncio.run(main())

output 是：

Produced 0
Consumed 0
Produced 1
Consumed 1
Produced 2
Consumed 2
Produced 3
Cannot consume 3
Produced 4
Cannot consume 4
Produced 5
Produced 6
Produced 7
Produced 8
Produced 9
Task exception was never retrieved
future: <Task finished name='Task-3' coro=<consume() done, defined at test.py:9> exception=ValueError('3 too big')>
Traceback (most recent call last):
  File "test.py", line 14, in consume
    raise ValueError(f'{task} too big')
ValueError: 3 too big
Task exception was never retrieved
future: <Task finished name='Task-2' coro=<consume() done, defined at test.py:9> exception=ValueError('4 too big')>
Traceback (most recent call last):
  File "test.py", line 14, in consume
    raise ValueError(f'{task} too big')
ValueError: 4 too big

有沒有辦法在消費者提出異常后通知生產者停止生產？
上面的代碼使用了多個生產者。 如果“通知”機制只能在單生產者模式下工作，也是可以接受的。

Answer 1

靈感來自user4815162342的建議

設置全局 Boolean 變量以防消費者異常，並在生產者中檢查

import asyncio

stop = False


async def single_produce(q: asyncio.Queue):
    global stop
    for task in range(10):
        await asyncio.sleep(0.001)
        if stop:
            break
        await q.put(task)
        print(f'Produced {task}')


async def multi_produce(q: asyncio.Queue, task):
    await asyncio.sleep(0.001)
    await q.put(task)
    print(f'Produced {task}')


async def consume(q: asyncio.Queue):
    global stop
    while True:
        task = await q.get()
        if task > 2:
            stop = True
            print(f'Cannot consume {task}')
            raise ValueError(f'{task} too big')
        print(f'Consumed {task}')
        q.task_done()


async def main(mode):
    global stop
    queue = asyncio.Queue(1)
    consumers = [asyncio.create_task(consume(queue)) for _ in range(2)]
    if mode == 'single':
        print('single producer')
        await asyncio.create_task(single_produce(queue))
    elif mode == 'multiple':
        print('multiple producers')
        for i in range(10):
            if stop:
                break
            await asyncio.create_task(multi_produce(queue, i))
    await asyncio.wait([queue.join(), *consumers],
                       return_when=asyncio.FIRST_COMPLETED)


asyncio.run(main('single'))
# asyncio.run(main('multiple'))