From 71ade7309218a193ab7733fb79bd207e2fc49d1f Mon Sep 17 00:00:00 2001 From: Douglas Raillard Date: Tue, 7 May 2024 10:20:02 +0100 Subject: [PATCH 1/3] target: Provide context manager API for Target Allow cleanly disconnecting the Target object, so that we don't get garbage output from __del__ later on when half of the namespace has already disappeared. --- devlib/target.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/devlib/target.py b/devlib/target.py index 5f2c595a0..a2d0562c4 100644 --- a/devlib/target.py +++ b/devlib/target.py @@ -526,6 +526,18 @@ def disconnect(self): if self._async_pool is not None: self._async_pool.__exit__(None, None, None) + def __enter__(self): + return self + + def __exit__(self, *args, **kwargs): + self.disconnect() + + async def __aenter__(self): + return self.__enter__() + + async def __aexit__(self, *args, **kwargs): + return self.__exit__(*args, **kwargs) + def get_connection(self, timeout=None): if self.conn_cls is None: raise ValueError('Connection class not specified on Target creation.') From 83c41a07c4975da1a88ddf56a044a34963c316a2 Mon Sep 17 00:00:00 2001 From: Douglas Raillard Date: Tue, 7 May 2024 10:27:50 +0100 Subject: [PATCH 2/3] target: Make Target.disconnect() steal current connections Ensure the connections that Target.disconnect() closes do not stay around in case the Target object is later reused. --- devlib/target.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/devlib/target.py b/devlib/target.py index a2d0562c4..b8e5aad3e 100644 --- a/devlib/target.py +++ b/devlib/target.py @@ -521,6 +521,13 @@ async def check_connection(self): def disconnect(self): connections = self._conn.get_all_values() + # Now that we have all the connection objects, we simply reset the TLS + # property so that the connections we got will not be reused anywhere. + del self._conn + + unused_conns = self._unused_conns + self._unused_conns.clear() + for conn in itertools.chain(connections, self._unused_conns): conn.close() if self._async_pool is not None: From 319b697c4535c43d987d179d71fd570ce65d7941 Mon Sep 17 00:00:00 2001 From: Douglas Raillard Date: Tue, 7 May 2024 11:01:05 +0100 Subject: [PATCH 3/3] target: Run Target.disconnect() upon process termination Use atexit handler to run Target.disconnect() when the process is about to exit. This avoids running it with a half torn down namespace, with ensuing exceptions and non-clean disconnect. --- devlib/target.py | 13 +++++++++++-- devlib/utils/ssh.py | 7 ------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/devlib/target.py b/devlib/target.py index b8e5aad3e..907fe2807 100644 --- a/devlib/target.py +++ b/devlib/target.py @@ -13,6 +13,7 @@ # limitations under the License. # +import atexit import asyncio from contextlib import contextmanager import io @@ -40,6 +41,7 @@ from past.types import basestring from numbers import Number from shlex import quote +from weakref import WeakMethod try: from collections.abc import Mapping except ImportError: @@ -413,6 +415,10 @@ def kind_conflict(kind, names): )) self._modules = modules + atexit.register( + WeakMethod(self.disconnect, atexit.unregister) + ) + self._update_modules('early') if connect: self.connect(max_async=max_async) @@ -530,8 +536,11 @@ def disconnect(self): for conn in itertools.chain(connections, self._unused_conns): conn.close() - if self._async_pool is not None: - self._async_pool.__exit__(None, None, None) + + pool = self._async_pool + self._async_pool = None + if pool is not None: + pool.__exit__(None, None, None) def __enter__(self): return self diff --git a/devlib/utils/ssh.py b/devlib/utils/ssh.py index 0eb7db2ba..aa1e873bb 100644 --- a/devlib/utils/ssh.py +++ b/devlib/utils/ssh.py @@ -24,14 +24,12 @@ import socket import sys import time -import atexit import contextlib import select import copy import functools import shutil from shlex import quote -from weakref import WeakMethod from paramiko.client import SSHClient, AutoAddPolicy, RejectPolicy import paramiko.ssh_exception @@ -372,8 +370,6 @@ def __init__(self, self.client = None try: self.client = self._make_client() - 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". @@ -815,9 +811,6 @@ def __init__(self, self.conn = telnet_get_shell(host, username, password, port, timeout, original_prompt) - weak_close = WeakMethod(self.close, atexit.unregister) - atexit.register(weak_close) - def fmt_remote_path(self, path): return '{}@{}:{}'.format(self.username, self.host, path)