sascha/devlib
1
0
Fork 0
mirror of https://github.com/ARM-software/devlib.git synced 2025-01-30 17:50:46 +00:00
Code Issues Releases Wiki Activity

utils/ssh: Fix atexit.register() SshConnection leak

Browse Source 
SshConnection registers an atexit handler so the connection is closed
upon exiting the process if it has not been done before. However, the
handler keeps a reference on the connection, which means it _will_ stay
alive. If lots of short-lived connections are created (which can happen
when using e.g. ThreadPoolExecutor), they will simply stay around and
leak.

Fix that by using a weak reference (WeakMethod) to register in the
atexit handler, with a callback to unregister it when the object is
deallocated.
This commit is contained in:
Douglas Raillard 2024-03-28 14:47:24 +00:00 committed by Marc Bonnici
parent 5817866ad0
commit 28b30649f1
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
devlib/utils/ssh.py
View File

@ -30,6 +30,7 @@ import select
import copy
import functools
from shlex import quote
from weakref import WeakMethod
from paramiko.client import SSHClient, AutoAddPolicy, RejectPolicy
import paramiko.ssh_exception
@ -370,7 +371,8 @@ class SshConnection(SshConnectionBase):
self.client = None
try:
self.client = self._make_client()
atexit.register(self.close)
weak_close = WeakMethod(self.close, atexit.unregister)
atexit.register(weak_close)
# Use a marker in the output so that we will be able to differentiate
# target connection issues with "password needed".
@ -811,7 +813,9 @@ class TelnetConnection(SshConnectionBase):
timeout = timeout if timeout is not None else self.default_timeout
self.conn = telnet_get_shell(host, username, password, port, timeout, original_prompt)
atexit.register(self.close)
weak_close = WeakMethod(self.close, atexit.unregister)
atexit.register(weak_close)
def fmt_remote_path(self, path):
return '{}@{}:{}'.format(self.username, self.host, path)