From 28b30649f136c498df1a7ddf4b56af78df821331 Mon Sep 17 00:00:00 2001 From: Douglas Raillard Date: Thu, 28 Mar 2024 14:47:24 +0000 Subject: [PATCH] utils/ssh: Fix atexit.register() SshConnection leak 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. --- devlib/utils/ssh.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/devlib/utils/ssh.py b/devlib/utils/ssh.py index a998b26..fbff5ca 100644 --- a/devlib/utils/ssh.py +++ b/devlib/utils/ssh.py @@ -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)