Simultaneous cancellation of nurseries in trio

trio_cancel_services.py

"""
This tests shows a trio behavior that we did not anticipated.

When a cancellation is issued in trio, all the nurseries in the scope
are cancelled simultaneously. This can sometimes lead to confusing
errors, especially when those nurseries are hidden in context
managers.

This example shows the case of two services, implemented as an async
context manager running a background task. Then, we assume that the
second service depends on the first one, up until the end of its
teardown.

As demonstrated by the test, the nesting of the two services works
fine in the nominal case or when an exception is raised, but fails
when the main scope is cancelled as both services get cancelled
simultaneously.
"""

import trio
import pytest
from contextlib import asynccontextmanager


@asynccontextmanager
async def run_service(depends_on=None):

    class service:
        pass

    async def target():
        try:
            service.running = True
            await trio.sleep_forever()
        finally:
            service.running = False
            with trio.CancelScope(shield=True):
                await trio.sleep(.1)
            assert depends_on is None or depends_on.running

    async with trio.open_nursery() as nursery:
        nursery.start_soon(target)
        yield service
        nursery.cancel_scope.cancel()


@pytest.mark.trio
async def test_simultaneous_cancellation():
    # Works fine in nominal case
    async with run_service() as service1:
        async with run_service(depends_on=service1):
            await trio.sleep(.1)

    # Works fine when an exception is raised
    with pytest.raises(ZeroDivisionError):
        async with run_service() as service1:
            async with run_service(depends_on=service1):
                await trio.sleep(.1)
                1/0

    # Does not work when the outer scope is cancelled
    with trio.CancelScope() as scope:
        async with run_service() as service1:
            async with run_service(depends_on=service1):
                await trio.sleep(.1)
                scope.cancel()