How to pipeline data from one stream to another with multiple consumers

[geospatial-jeff]

Hello! Thanks for the great library!! I'm trying to pipeline data from one stream to another to implement a simple fan-out / map-reduce pattern. The data flow is quite simple:

1. A producer sends data to the first stream, in this case integers for range(10).
2. Multiple consumers process the stream, putting the output on a new stream (map).
3. A second consumer processes the final stream (reduce).

producer -> stream -> map_consumer(s) -> stream -> reduce_consumer

This seems to work fine when I have a single consumer, for example:

import anyio

async def map_consumer(receive1, send2):
    async with send2:
        async with receive1:
            async for item in receive1:
                print("map consumer - ", item)
                await send2.send(item)
        print("closed receive1")
    print("closed send2")


async def reduce_consumer(receive2):
    async with receive2:
        async for item in receive2:
            print("reduce consumer - ", item)
    print("closed receive2")


async def producer(send1):
    async with send1:
        for idx in range(10):
            print("producing - ", idx)
            await send1.send(idx)
    print("closed send1")


async def main():
    send1, receive1 = anyio.create_memory_object_stream(max_buffer_size=10)
    send2, receive2 = anyio.create_memory_object_stream(max_buffer_size=10)

    async with anyio.create_task_group() as tg:
        # start the producer
        tg.start_soon(producer, send1)

        # start the consumer
        tg.start_soon(map_consumer, receive1, send2)

        # start a second consumer
        tg.start_soon(reduce_consumer, receive2)


anyio.run(main)

However when I add multiple map consumers, using clone to fork the stream across each consumer, the data processes as expected but the task group never exists. The print statements in the script indicates that streams send1, receive1, and send2 are all closing, however receive2 never closes. The receive2 stream is wrapped in an async context manager, so I'm a bit confused why it never closes.

async def main():
    send1, receive1 = anyio.create_memory_object_stream(max_buffer_size=10)
    send2, receive2 = anyio.create_memory_object_stream(max_buffer_size=10)

    async with anyio.create_task_group() as tg:
        # start the producer
        tg.start_soon(producer, send1)

        # start the consumer
        for _ in range(2):
            tg.start_soon(map_consumer, receive1.clone(), send2.clone())

        # start a second consumer
        tg.start_soon(reduce_consumer, receive2)

What am I doing wrong? Thanks for the help!

[smurfix]

That's easy. You need to close the original streams, i.e. receive1 and send2.

[geospatial-jeff]

awesome, thanks!

async def main():
    send1, receive1 = anyio.create_memory_object_stream(max_buffer_size=10)
    send2, receive2 = anyio.create_memory_object_stream(max_buffer_size=10)

    async with anyio.create_task_group() as tg:
        # start the producer
        tg.start_soon(producer, send1)

        # start the consumer
        consumers = 2
        streams = [(receive1, send2)]
        for _ in range(consumers - 1):
            streams.append((receive1.clone(), send2.clone()))
        
        for stream in streams:
            tg.start_soon(map_consumer, *stream)

        # start a second consumer
        tg.start_soon(reduce_consumer, receive2)