import aiosched



async def producer(c: aiosched.NetworkChannel, n: int):
    print("[producer] Started")
    for i in range(n):
        await c.write(str(i).encode())
        print(f"[producer] Sent {i}")
        await aiosched.sleep(0.5)   # This makes the receiver wait on us!
    await c.close()
    print("[producer] Done")


async def consumer(c: aiosched.NetworkChannel):
    print("[receiver] Started")
    try:
        while await c.pending():
            item = await c.read(1)
            print(f"[consumer] Received {item.decode()}")
            # await aiosched.sleep(2)    # If you uncomment this, the except block will be triggered
    except aiosched.errors.ResourceClosed:
        print("[consumer] Stream has been closed early!")
    print("[consumer] Done")


async def main(channel: aiosched.NetworkChannel, n: int):
    t = aiosched.clock()
    print("[main] Starting children")
    async with aiosched.with_context() as ctx:
        await ctx.spawn(consumer, channel)
        await ctx.spawn(producer, channel, n)
    print(f"[main] All done in {aiosched.clock() - t:.2f} seconds")


aiosched.run(
    main, aiosched.NetworkChannel(), 5, debugger=()
)  # 2 is the max size of the channel