A bad service manager in trio

bad_service_manager.py

import trio
from contextlib import AsyncExitStack, asynccontextmanager


@asynccontextmanager
async def some_service():

    class Service():
        def __init__(self, nursery):
            self.nursery = nursery

        async def do_something(self):
            self.nursery.start_soon(trio.sleep, 1)
            await trio.sleep(.1)
            return 3

    async with trio.open_nursery() as nursery:
        yield Service(nursery)


@asynccontextmanager
async def service_manager():

    class ServiceManager():
        def __init__(self, stack):
            self.stack = stack
            self.cache = {}

        async def get_service(self, n):
            if n not in self.cache:
                self.cache[n] = await stack.enter_async_context(some_service())
            return self.cache[n]

    async with AsyncExitStack() as stack:
        yield ServiceManager(stack)


async def main():
    async with trio.open_nursery() as global_nursery:

        async with service_manager() as manager:

            # Do something with service 1
            s1 = await manager.get_service(1)
            assert await s1.do_something() == 3

            async def target(manager):
                # OK because service 1 has already been started
                s1 = await manager.get_service(1)
                assert await s1.do_something() == 3

                # Not OK, this will produce the messy bug
                s2 = await manager.get_service(2)
                assert await s2.do_something() == 3

            async with trio.open_nursery() as local_nursery:

                # It doesn't matter which nursery is used here
                (global_nursery or local_nursery).start_soon(target, manager)
                await trio.sleep(.1)


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