vxgmichel / aioconsole

Asynchronous console and interfaces for asyncio
http://aioconsole.readthedocs.io
GNU General Public License v3.0
457 stars 39 forks source link

Termination of `AsynchronousCli` task with task groups #104

Closed aunsbjerg closed 1 year ago

aunsbjerg commented 1 year ago

I'm trying to use aioconsole and AsynchronousCli in my asyncio based application. In my application, I use task groups for creating tasks, including my AsynchronousCli-based cli. Something like:

async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(AsynchronousCli(commands).interact(), name="cli")

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

Now, the command line part is working just fine, no problems. But I cannot figure out how to properly terminate the application. Without aioconsole, asyncio.run(main()) install sigint handlers so that ctrl+c terminates the program - as expected. But aioconsole captures ctrl+c and doesn't terminate. If I use the built-in exit command, the application terminates but with an exception regarding Task exception was never retrieved.

I'm not sure, but it seems to me like the problem is in the way the _interact() of the AsynchronousConsole class is dealing with asyncio.CancelledError, e.g. just ignoring it completely. For my needs, the following patch works:

diff --git a/aioconsole/console.py b/aioconsole/console.py
index ae75917..538a1bb 100644
--- a/aioconsole/console.py
+++ b/aioconsole/console.py
@@ -204,6 +204,7 @@ class AsynchronousConsole(code.InteractiveConsole):
                 await self.flush()
                 self.resetbuffer()
                 more = 0
+                break

     async def push(self, line):
         self.buffer.append(line)

But the unit tests now fail so I'm not sure if that is the right approach.

vxgmichel commented 1 year ago

Interesting, thanks for the report!

First of all, note that interact has two arguments that can help with your case:

However there's still a problem with the code as it is: it does not discriminate between CancelledError that originates from the handler that aioconsole set up, and a CancelledError that might originate from other parts of the code like TaskGroup for instance. The following patch should fix the issue:

diff --git a/aioconsole/console.py b/aioconsole/console.py
index ae75917..b88d0ca 100644
--- a/aioconsole/console.py
+++ b/aioconsole/console.py
@@ -63,6 +63,8 @@ class AsynchronousConsole(code.InteractiveConsole):
         self.locals["print"] = self.print
         self.locals["help"] = self.help
         self.locals["ainput"] = self.ainput
+        # Internals
+        self._sigint_received = False

     @functools.wraps(print)
     def print(self, *args, **kwargs):
@@ -120,6 +122,7 @@ class AsynchronousConsole(code.InteractiveConsole):
         self.buffer = []

     def handle_sigint(self, task):
+        self._sigint_received = True
         task.cancel()
         if task._fut_waiter._loop is not self.loop:
             task._wakeup(task._fut_waiter)
@@ -200,6 +203,11 @@ class AsynchronousConsole(code.InteractiveConsole):
                 else:
                     more = await self.push(line)
             except asyncio.CancelledError:
+                # Not our cancellation
+                if not self._sigint_received:
+                    raise
+                # Manage cancellation
+                self._sigint_received = False
                 self.write("\nKeyboardInterrupt\n")
                 await self.flush()
                 self.resetbuffer()

I'll add this patch to an upcoming PR. Does that help?

aunsbjerg commented 1 year ago

Thanks for the thorough response, much appreciated. I have applied your patch and done some testing with various combinations of stop and handle_sigint:

So at least for my application, stop=False and handle_sigint=True is the closest to what I need. I would still like to propagate the KeyboardInterrupt to close the application on exit, so I found that adding a done callback does the trick:

def doneCallback(_):
    raise KeyboardInterrupt

taskGroup.create_task(
    AsynchronousCli(commands).interact(
        stop=False,
        handle_sigint=True,
    ),
    name="cli",
).add_done_callback(doneCallback)

Possibly, that could be added into AsynchronousCli or it's parents.

vxgmichel commented 1 year ago

Nice report! That juste gave me an idea, what about applying this patch and using stop=True and handle_sigint=True?

diff --git a/aioconsole/console.py b/aioconsole/console.py
index b88d0ca..199671f 100644
--- a/aioconsole/console.py
+++ b/aioconsole/console.py
@@ -158,6 +158,8 @@ class AsynchronousConsole(code.InteractiveConsole):
             if handle_sigint:
                 self.add_sigint_handler()
             await self._interact(banner)
+            if stop:
+                raise SystemExit
         # Exit
         except SystemExit:
             if stop:
@@ -166,8 +168,6 @@ class AsynchronousConsole(code.InteractiveConsole):
         finally:
             if handle_sigint:
                 self.remove_sigint_handler()
-            if stop:
-                self.loop.stop()

     async def _interact(self, banner=None):
         # Get ps1 and ps2

Would this be the behavior you're looking for? I'd like to get rid of this loop.stop() anyway.

aunsbjerg commented 1 year ago

With both patches applied and stop=True + handle_sigint=True, I get Task exception was never retrieved on exit and ctrl+c does not terminate.

The "problem" is in the way SystemExit is dealt with by asyncio (as far as I can tell at least, I'm no expert). The task group context manager will happily propagate SystemExit, but for some reason a SystemExit does not cause the main() task to get cancelled properly. In contrast, KeyboardInterrupt does cancel main().

So basically, if the CLI raises SystemExit, I'll need to wrap the task group context manager in a try block to catch the SystemExit error, like so:

async def main():
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(AsynchronousCli(commands).interact(), name="cli")

    except SystemExit:
        pass

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

In that case, I think I would prefer the CLI task just raising asyncio.CancelledError and let the application deal with when and how to exit.

vxgmichel commented 1 year ago

The "problem" is in the way SystemExit is dealt with by asyncio (as far as I can tell at least, I'm no expert). The task group context manager will happily propagate SystemExit, but for some reason a SystemExit does not cause the main() task to get cancelled properly. In contrast, KeyboardInterrupt does cancel main().

I investigated a bit more and I don't fully agree with what you described. From what I can tell SystemExit and KeyboardInterrupt work in a similar way. The difference between the patch (adding if stop: raise SystemExit) and the add_done_callback approach you mentioned earlier is that in the first case, the exception is raised in the subtask and in the second case, it's raised in the main task. Note that in both cases it stops the application (i.e the exception bubbled up correctly). The difference is that when the SystemExit/KeyboardInterrupt exception is raised in the subtask, asyncio will show a warning prompting that the Task exception was never retrieved, along with the corresponding traceback. This is probably because those BaseExceptions end up stopping the loop without going through the proper task shutdown logic.

I would argue that this is a bug in asyncio since the documentation clearly states that those exceptions should be handled properly:

Two base exceptions are treated specially: If any task fails with KeyboardInterrupt or SystemExit, the task group still cancels the remaining tasks and waits for them, but then the initial KeyboardInterrupt or SystemExit is re-raised instead of ExceptionGroup or BaseExceptionGroup.

I'll report it on the python tracker, since it can be reproduced very easily:

async def raise_system_exit():
    raise SystemExit

async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(raise_system_exit())

In the meantime, I think the workaround you found is fine:

async def main():
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(AsynchronousCli(commands).interact(), name="cli")

    except SystemExit:
        pass

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

In that case, I think I would prefer the CLI task just raising asyncio.CancelledError and let the application deal with when and how to exit.

I'm not sure we want to do that, since asyncio.CancelledError has a specific meaning in the context of asyncio. I'd argue the following logic is the way to go:

class StopTaskGroup(Exception):
    pass

async def cli():
    await AsynchronousCli(commands).interact(stop=False)
    # The CLI has exited, do something about it
    ...
    # Maybe cancel the current task group
    raise StopTaskGroup

async def main():
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(cli(), name="cli")
            while True:
                await asyncio.sleep(1)
    # The CLI has exited and cancelled the task group
    except* StopTaskGroup:
        # Maybe do something else here
        ...

Hope that makes sense, thank you for your valuable feedback :)

aunsbjerg commented 1 year ago

I would argue that this is a bug in asyncio since the documentation clearly states that those exceptions should be handled properly:

Two base exceptions are treated specially: If any task fails with KeyboardInterrupt or SystemExit, the task group still cancels the remaining tasks and waits for them, but then the initial KeyboardInterrupt or SystemExit is re-raised instead of ExceptionGroup or BaseExceptionGroup.

I agree, I was wondering why SystemExit was treated differently that KeyboardInterrupt given the documentation, but that makes sense if there is a bug.

With the solutions and workarounds we have discussed, I am able to proceed so feel free to close this issue once you are ready - and thanks for the help! :)

vxgmichel commented 1 year ago

With the solutions and workarounds we have discussed, I am able to proceed so feel free to close this issue once you are ready

Great, I went ahead and merged #105. It'll be available in the next release :)

Also I reported the asyncio issue here: https://github.com/python/cpython/issues/101515

vxgmichel commented 1 year ago

Version v0.6.0 is out :tada: