trio nursery that doesn't cancel all tasks if one fails

Question

I'd like to implement a server with trio. Individual client connections are handled by tasks spawned by a nursery. However, the trio docs say that "If any task inside the nursery finishes with an unhandled exception, then the nursery immediately cancels all the tasks inside the nursery.". This is pretty unfortunate for my use-case. I'd much rather continue servering the other connections while logging the error. Is there a way to do that?

Answer 1

You can make your own implementation of the nursery interface:

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

    @property
    def cancel_scope(self):
        return self.nursery.cancel_scope

    async def _run_and_log_errors(self, async_fn, *args):
        # This is more cumbersome than it should be
        # See https://github.com/python-trio/trio/issues/408
        def handler(exc):
            if not isinstance(exc, Exception):
                return exc
            logger.error("Unhandled exception!", exc_info=exc)
        with trio.MultiError.catch(handler):
            return await async_fn(*args)

    def start_soon(self, async_fn, *args, **kwargs):
        self.nursery.start_soon(self._run_and_log_errors, async_fn, *args, **kwargs)

    async def start(self, async_fn, *args, **kwargs):
        return await self.nursery.start(self._run_and_log_errors, async_fn, *args, **kwargs)

@asynccontextmanager
async def open_exception_logging_nursery():
    async with trio.open_nursery() as nursery:
        yield ExceptionLoggingNursery(nursery)

Note that we only catch Exception subclasses, and allow other exceptions to continue propagate. This means that if one of your child tasks raises a KeyboardInterrupt (because you hit control-C), or a trio.Cancelled (because you, well, cancelled it... maybe because you hit control-C and the nursery that the parent is living in got cancelled), then those exceptions are allowed to propagate out and still cause all of the other tasks to be cancelled, which is almost certainly what you want.

It's a bit of code, but it can easily be put in a reusable library. (If I were doing this for real I might make the exception-handling code an argument passed to open_exception_logging_nursery, instead of hard-coding that call to logger.error.) And I'd love to see a library with this kind of "smart supervisor" in it -- the basic trio nursery was always intended as a building block for such things. You can imagine other even more interesting policies too, like "if a task exits with an unhandled exception, log something and then restart it, with an exponential backoff". (Erlang supervisors are a good source of ideas to steal be inspired by.)