sascha/devlib
1
0
Fork 0
mirror of https://github.com/ARM-software/devlib.git synced 2025-01-31 02:00:45 +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 copy
import functools import functools
from shlex import quote from shlex import quote
from weakref import WeakMethod
from paramiko.client import SSHClient, AutoAddPolicy, RejectPolicy from paramiko.client import SSHClient, AutoAddPolicy, RejectPolicy
import paramiko.ssh_exception import paramiko.ssh_exception
@ -370,7 +371,8 @@ class SshConnection(SshConnectionBase):
self.client = None self.client = None
try: try:
self.client = self._make_client() 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 # Use a marker in the output so that we will be able to differentiate
# target connection issues with "password needed". # target connection issues with "password needed".
@ -811,7 +813,9 @@ class TelnetConnection(SshConnectionBase):
timeout = timeout if timeout is not None else self.default_timeout timeout = timeout if timeout is not None else self.default_timeout
self.conn = telnet_get_shell(host, username, password, port, timeout, original_prompt) 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): def fmt_remote_path(self, path):
return '{}@{}:{}'.format(self.username, self.host, path) return '{}@{}:{}'.format(self.username, self.host, path)