Graceful termination when using ProxyCommand

[slietar]

Hi,

First of all thank you for writing this library, it has been incredibly useful to me so far.

I am currently running into a few issues while connecting to a server using the ProxyCommand option:

1. When sending a SIGINT to the Python process, the signal is also received by the proxy command process and causes it to terminate. This prevents any further communication and thus the graceful termination of the Python program. Here is a short example to reproduce this; pressing Ctrl+C leads to an unexpected "SSH connection closed" exception.

   async with asyncssh.connect("...") as conn:
       print("Connected")
   
       try:
           await asyncio.Future()
       finally:
           await conn.run("echo Hello")

   I believe this can be fixed quite easily by setting start_new_session=True on loop.subprocess_exec to prevent signals from propagating.
   

   asyncssh/asyncssh/connection.py

   Line 410 in c7bc3aa

   _, tunnel = await loop.subprocess_exec(_ProxyCommandTunnel, *command)

2. The context manager of SSHConnection does not wait for the proxy command process to terminate before exiting. This causes a leak of the process which is reported as a warning when the event loop closes: "Loop [...] that handles pid xxx is closed". It should be possible to wait for process termination by adding a process_exited method in _ProxyCommandTunnel, but I am unsure on how to proceed from there as the connection cleanup logic seems quite complex.

[ronf]

Thanks so much for the report!

Here's a potential fix, based on your suggestions:

diff --git a/asyncssh/connection.py b/asyncssh/connection.py
index 82a5ec1..7e4c9c7 100644
--- a/asyncssh/connection.py
+++ b/asyncssh/connection.py
@@ -387,6 +387,11 @@ async def _open_proxy(

             self._conn.connection_lost(exc)

+        def process_exited(self):
+            """Called when the child process has exited"""
+
+            self._close_event.set()
+
         def write(self, data: bytes) -> None:
             """Write data to this tunnel"""

@@ -402,12 +407,14 @@ async def _open_proxy(
             """Close this tunnel"""

             if self._transport: # pragma: no cover
-                self._transport.close()
+                self._conn.create_task(self._close_event.wait())

-            self._close_event.set()
+                self._transport.close()
+                self._transport = None


-    _, tunnel = await loop.subprocess_exec(_ProxyCommandTunnel, *command)
+    _, tunnel = await loop.subprocess_exec(_ProxyCommandTunnel, *command,
+                                           start_new_session=True)

     return cast(_Conn, cast(_ProxyCommandTunnel, tunnel).get_conn())

In my limited testing here, this avoids both of the issues you raised, and hitting Ctrl-C no longer generates errors about the event loop already being closed. I didn't see the error about the pid not being closed before this change, but I did see other errors such as ignored exceptions. This all goes away for me with the above change.

If this works for you, I'll check it into the develop branch.

[slietar]

Thanks for your reply. Unfortunately I still occasionally see the warning; whether it appears or not might depend on how fast the process terminates so I'm not surprised that it's not always visible.

I think the remaining problem is caused by the fact that asyncio.run cancels all remaining tasks when stopping, meaning self._close_event.wait() is cancelled and since it was not awaited anywhere else, we never wait for the process to terminate. A possible fix would be to wait for all tasks of the SSHConnection to finish when exiting, i.e. adding asyncio.wait(self._tasks) somewhere in SSHConnection.__aexit__. This also makes more sense to me as a user as I would expect everything to be finished once the context manager exits. I've tried it and it seems to work, but it may have unintended consequences so let me know what you think.

[ronf]

Are you explicitly stopping or closing the event loop? As far as I know, asyncio.run() should only exit when all tasks scheduled on the event loop it created have completed. Applications may speed up that process by explicitly calling Task.cancel() (and this is particularly important for "daemon" tasks that might not otherwise complete on their own), but unless something explicitly stops the event loop, I wouldn't expect any automatic cancellation to occur here.

One of the challenges is that not all callers use context managers or explicitly call wait_closed() on SSHConnections. So, it needs to make sure it starts the process of cleaning everything up when close() or abort() is called on it, but those methods are not asynchronous and can't actually wait for anything.

It would be possible to add a wait_closed() method to the _ProxyCommandTunnel class, but then I'd need to find a way to track that object in SSHConnection so its _cleanup() can wait for the process to exit. This is messy, though, as SSHConnection usually only sees a transport object, and this is a protocol object we'd now need to expose. The _ProxyCommandTunnel object does have access to the SSHConnection, so it could call a new SSHConnection method to set itself as something else that needs to be waited on, but having it schedule a queued task to wait for the process to exit really should be run before exiting the event loop.

[ronf]

I did a little more digging on this. It looks like the behavior may have changed somewhat in Python 3.11. The SIGINT handler by default will cancel the main task, but once it exits, any other remaining tasks are apparently cancelled. That may be what you were seeing.

I'm not sure if exiting the main task always cancels the other tasks or if this only happens when the SIGINT handler runs. If the latter is the case, you might be able to solve this issue by catching KeyboardInterrupt yourself and having that jump you out of the SSHConnection's context manager. If this doesn't cancel other tasks, that'll allow a graceful cleanup, including running the task which waits for process_exited().

[ronf]

After some additional testing, I confirmed that catching asyncio.CancelledError from within the main task is enough to prevent other tasks from being cancelled.

That said, without some kind of explicit wait fall in that except block, there's no guarantee that all the remaining tasks will complete before the event loop closes, so you may not get to the process exit before then. However, I haven't been able to get my test code to print any errors during the event loop shutdown. So, you may want to give this a try and see if it helps.

[slietar]

My understanding is that when the user presses Ctrl+C this causes of cancellation of the root task that asyncio.run is running (see here). Once that task exits, whether it is successfully or not, all other tasks are cancelled (see here) and awaited without any additional user action. But if using structured concurrency, which is more or less implied by the use of a context manager, I think there shouldn't be any tasks left anyway.

If I understood correctly, you're suggesting to await remaining tasks manually after exiting the context manager, something like this:

try:
    async with asyncssh.connect("...") as conn:
        ...
finally:
    await asyncio.gather(*(asyncio.all_tasks() - {asyncio.current_task()}))

This works but makes assumptions on how the user organizes their tasks since it's impossible to differentiate tasks created by asyncssh from other tasks. An exception raised by an asyncssh task would be reported with an obscure traceback rather than one that points to the call to asyncssh.connect.

I initially tried to implement something like what you mention with wait_closed() but ran into the same problems as you, which is why was hoping you would have a better idea on how to fix this haha. Having something like wait_closed() that's called when exiting the context manager but can also be called manually by users who use a synchronous close() seems like the best approach – but as you said, the current design makes it tricky to implement.

[ronf]

Thanks for the pointers to the cancel handling. I wasn't suggesting trying to wait on all_tasks(), as that would include tasks which have nothing to do with the task related to this subprocess. It may even include tasks not related to AsyncSSH, and there's no guarantee any of those tasks will finish (especially if they're not cancelled first). Anything we do here would have to be specific to the _ProxyCommandTunnel instance and its associated subprocess.

The good news here is that the tunnel argument returned by loop.subprocess_exec is actually the _ProxyCommandTunnel instance (the "protocol") that wait_closed() could be added to. However, _open_proxy would need to be changed to return that. Right now, it only returns the associated connection object. That's doable, though, and there's even already a member of SSHConnection that tracks the underlying tunnel object. Right now it is only used when tunneling over another SSH connection opened in _open_tunnel, and it's used to trigger a close() on a tunnel object that was opened without the caller's involvement. It seems like that internal tunnel could have a similar issue where wait_closed() is not called on it and it doesn't get to fully clean up if the main task exits. So, perhaps I can provide a way for both types of tunnels to be waited on when closing the SSH connection within the tunnel. I'll look into this.

[ronf]

I figured out how to get wait_closed() on a _ProxyCommand_Tunnel instance to be called when the SSHConnection's wait_closed() is called, which is already called from the SSHConnection's context manager. However, when I hit Ctrl-C while sitting inside the connection's context manager, I don't actually see wait_closed() being called on the SSHConnection, and adding a try block around that doesn't seem to make a difference. The context manager exit code is never called when Ctrl-C is hit, but works just fine if the process is left to exit normally (like after running a sleep command). I thought I had this working before, but I'm not sure what has changed.

[ronf]

Here's the latest diff I've been experimenting with:

diff --git a/asyncssh/connection.py b/asyncssh/connection.py
index e81a92c..aa56b59 100644
--- a/asyncssh/connection.py
+++ b/asyncssh/connection.py
@@ -196,6 +196,9 @@ class _TunnelProtocol(Protocol):
     def close(self) -> None:
         """Close this tunnel"""

+    async def wait_closed(self):
+        """Wait for this tunnel to close"""
+
 class _TunnelConnectorProtocol(_TunnelProtocol, Protocol):
     """Protocol to open a connection to tunnel an SSH connection over"""

@@ -387,6 +390,11 @@ async def _open_proxy(

             self._conn.connection_lost(exc)

+        def process_exited(self):
+            """Called when the child process has exited"""
+
+            self._close_event.set()
+
         def write(self, data: bytes) -> None:
             """Write data to this tunnel"""

@@ -403,13 +411,20 @@ async def _open_proxy(

             if self._transport: # pragma: no cover
                 self._transport.close()
+                self._transport = None

-            self._close_event.set()
+        async def wait_closed(self):
+            """Wait for this subprocess to exit"""
+
+            await self._close_event.wait()

+    _, tunnel = await loop.subprocess_exec(_ProxyCommandTunnel, *command,
+                                           start_new_session=True)

-    _, tunnel = await loop.subprocess_exec(_ProxyCommandTunnel, *command)
+    conn = cast(_Conn, cast(_ProxyCommandTunnel, tunnel).get_conn())
+    conn.set_tunnel(tunnel)

-    return cast(_Conn, cast(_ProxyCommandTunnel, tunnel).get_conn())
+    return conn


 async def _open_tunnel(tunnels: object, options: _Options,
@@ -1090,15 +1105,15 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):

             self._owner = None

+        if self._tunnel:
+            self._tunnel.close()
+            self._tunnel = None
+
         self._cancel_login_timer()
         self._close_event.set()

         self._inpbuf = b''

-        if self._tunnel:
-            self._tunnel.close()
-            self._tunnel = None
-
     def _cancel_login_timer(self) -> None:
         """Cancel the login timer"""

@@ -2851,6 +2866,9 @@ class SSHConnection(SSHPacketHandler, asyncio.Protocol):
         if self._agent:
             await self._agent.wait_closed()

+        if self._tunnel:
+            await self._tunnel.wait_closed()
+
         await self._close_event.wait()

     def disconnect(self, code: int, reason: str,

The proxy tunnel process is cleaned up nicely here in cases that don't involve a KeyboardInterrupt, but when a KeyboardInterrupt is triggered none of the code in AsyncSSH gets to run. I don't even seem to be able to get an except block for an asyncio.CancelledError exception to run from within the main task. It seems to just be canceling that task immediately on receiving the KeyboardInterrupt.

I can catch the KeyboardInterrupt, but only if I do it in the method calling loop.run_until_complete, outside of any async task.

Because tasks are cancelled so abruptly without seeming to unwind the stack, I also haven't been able to get things like process_exited to run on the proxy tunnel object when a KeyboardInterrupt is triggered, even though it works just fine during a normal close.

Any thoughts?

[ronf]

I spent some more time on this today and discovered something interesting...

My test code was using loop.run_until_complete() to run a run_client() function which included an "async with" on asyncssh.connect(), and the SSHConnection context manager aexit code never ran. However, when I changed to asyncio.run() on the exact same code, suddenly I was getting the exit code to run, including waiting for the proxy tunnel to get into process_exited() before the KeyboardInterrupt except block was run. So, perhaps the code above is good enough to solve this problem, as long as you are using asyncio.run(). When you get a chance, could you give this a try?

The original code made sure to call loop.close(), but that still wasn't enough to get the exit handlers to be called.

[ronf]

I've gone ahead and checked in a fix for this in commit f7b8baf in the "develop" branch.