Reuse SSH Client

Allow the installer to re-use the paramiko SSH Client instead of creating a new one every time, resulting in slightly lower install time.

Regression: Tested by installing a AIO-SX system, PASS

Story: 2005051
Task: 48502

Change-Id: I1fcb0df2a0a04a2983c7f7f32bf12f3351c0e8d9
Signed-off-by: Felipe Freire <felipe.freire@encora.com>
This commit is contained in:
Felipe Freire 2023-10-20 15:14:25 -03:00
parent c827cf81e9
commit 71318415e9
3 changed files with 120 additions and 93 deletions

View File

@ -0,0 +1 @@
from .ssh_exception import InvalidSSHConnection

View File

@ -0,0 +1,2 @@
class InvalidSSHConnection(Exception):
pass

View File

@ -34,11 +34,14 @@ from consts.networking import NICs, OAM, MGMT, Serial
from consts.timeout import HostTimeout
from consts import env
from exceptions import InvalidSSHConnection
from Parser import handle_args
# Global vars
V_BOX_OPTIONS = None
SSH_CONNECTIONS = {}
# Network
OAM_CONFIG = [getattr(OAM, attr) for attr in dir(OAM) if not attr.startswith('__')]
@ -972,6 +975,21 @@ def _connect_to_ssh(node='floating'):
LOG.error("#### Failed SSH connection\nError: %s", repr(exc))
raise
SSH_CONNECTIONS[node] = ssh
return ssh
def ssh_handler(node='floating'):
"""
Handles the SSH connection. Tries to retrieve a already existing connection.
If it doesn't exist, or isn't active, creates a new one.
"""
try:
ssh = SSH_CONNECTIONS[node]
ssh_transport = ssh.get_transport()
if ssh_transport is None or not ssh_transport.is_active():
raise InvalidSSHConnection()
except (InvalidSSHConnection, KeyError):
ssh = _connect_to_ssh(node)
return ssh
@ -1000,19 +1018,21 @@ def connect_to_ssh(node='floating'):
Returns: return code of decorated function
"""
try:
ssh = _connect_to_ssh(node)
kwargs['ssh_client'] = ssh
return func(*args, **kwargs)
finally:
if ssh:
ssh.close()
ssh = ssh_handler(node)
kwargs['ssh_client'] = ssh
return func(*args, **kwargs)
return connect_to_ssh_wrapper
return connect_to_ssh_decorator
def close_ssh_connections(ssh_clients: dict):
"""
Closes the connection with all current created SSH Clients.
"""
for _, client in ssh_clients.items():
if client is not None:
client.close()
def stage_test_success():
"""Prints a log message indicating the execution of a test stage."""
@ -1790,7 +1810,7 @@ def run_custom_script(script, timeout, console, mode):
LOG.info(" console mode: %s", console)
LOG.info(" user mode: %s", mode)
if console == 'ssh':
ssh_client = _connect_to_ssh()
ssh_client = ssh_handler()
# pylint: disable=W0703, C0103
_, __, return_code = run_ssh_cmd(ssh_client, f"./{script}", timeout=timeout, mode=mode)
if return_code != 0:
@ -2491,6 +2511,7 @@ def signal_handler():
"""
print('You pressed Ctrl+C!')
close_ssh_connections(SSH_CONNECTIONS)
kpi.print_kpi_metrics()
sys.exit(1)
@ -2522,99 +2543,102 @@ def log_heading_msg(msg, pattern='#', panel_size=20):
# pylint: disable=invalid-name
if __name__ == "__main__":
kpi.init_kpi_metrics()
signal.signal(signal.SIGINT, signal_handler)
try:
kpi.init_kpi_metrics()
signal.signal(signal.SIGINT, signal_handler)
load_config()
load_config()
if V_BOX_OPTIONS.list_stages:
print(f"Defined setups: {list(STAGES_CHAINS.keys())}")
if V_BOX_OPTIONS.setup_type and V_BOX_OPTIONS.setup_type in AVAILABLE_CHAINS:
AVAILABLE_CHAINS = [V_BOX_OPTIONS.setup_type]
for stg_chain in AVAILABLE_CHAINS:
stg_no = 1
print(f"Stages for setup on: {stg_chain}")
for stage in STAGES_CHAINS[stg_chain]:
print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP], stg_no))
stg_no += 1
print("Available stages that can be used for --custom-stages:")
for stage in AVAILABLE_STAGES:
print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP]))
sys.exit(0)
if V_BOX_OPTIONS.list_stages:
print(f"Defined setups: {list(STAGES_CHAINS.keys())}")
if V_BOX_OPTIONS.setup_type and V_BOX_OPTIONS.setup_type in AVAILABLE_CHAINS:
AVAILABLE_CHAINS = [V_BOX_OPTIONS.setup_type]
for stg_chain in AVAILABLE_CHAINS:
stg_no = 1
print(f"Stages for setup on: {stg_chain}")
for stage in STAGES_CHAINS[stg_chain]:
print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP], stg_no))
stg_no += 1
print("Available stages that can be used for --custom-stages:")
for stage in AVAILABLE_STAGES:
print(wrap_stage_help(stage, STAGE_CALLBACKS[stage][HELP]))
sys.exit(0)
init_logging(V_BOX_OPTIONS.labname, V_BOX_OPTIONS.logpath)
LOG.info("Logging to directory: %s", (get_log_dir() + "/"))
init_logging(V_BOX_OPTIONS.labname, V_BOX_OPTIONS.logpath)
LOG.info("Logging to directory: %s", (get_log_dir() + "/"))
LOG.info("Install manages: %s controllers, %s workers, %s storages.",
V_BOX_OPTIONS.controllers, V_BOX_OPTIONS.workers, V_BOX_OPTIONS.storages)
LOG.info("Install manages: %s controllers, %s workers, %s storages.",
V_BOX_OPTIONS.controllers, V_BOX_OPTIONS.workers, V_BOX_OPTIONS.storages)
# Setup stages to run based on config
install_stages = []
if V_BOX_OPTIONS.custom_stages:
# Custom stages
install_stages = V_BOX_OPTIONS.custom_stages.split(',')
for stage in install_stages:
invalid_stages = []
if stage not in AVAILABLE_STAGES:
invalid_stages.append(stage)
if invalid_stages:
LOG.warning("Following custom stages are not supported: %s.\n" \
"Choose from: %s", invalid_stages, AVAILABLE_STAGES)
sys.exit(1)
else:
# List all stages between 'from-stage' to 'to-stage'
stages = STAGES_CHAINS[V_BOX_OPTIONS.setup_type]
from_stg_index = 0
to_stg_index = None
if V_BOX_OPTIONS.from_stage:
if V_BOX_OPTIONS.from_stage == 'start':
from_stg_index = 0
else:
from_stg_index = stages.index(V_BOX_OPTIONS.from_stage)
if V_BOX_OPTIONS.to_stage:
if V_BOX_OPTIONS.from_stage == 'end':
to_stg_index = -1
else:
to_stg_index = stages.index(V_BOX_OPTIONS.to_stage) + 1
if to_stg_index is not None:
install_stages = stages[from_stg_index:to_stg_index]
# Setup stages to run based on config
install_stages = []
if V_BOX_OPTIONS.custom_stages:
# Custom stages
install_stages = V_BOX_OPTIONS.custom_stages.split(',')
for stage in install_stages:
invalid_stages = []
if stage not in AVAILABLE_STAGES:
invalid_stages.append(stage)
if invalid_stages:
LOG.warning("Following custom stages are not supported: %s.\n" \
"Choose from: %s", invalid_stages, AVAILABLE_STAGES)
sys.exit(1)
else:
install_stages = stages[from_stg_index:]
LOG.info("Executing %s stage(s): %s.", len(install_stages), install_stages)
# List all stages between 'from-stage' to 'to-stage'
stages = STAGES_CHAINS[V_BOX_OPTIONS.setup_type]
from_stg_index = 0
to_stg_index = None
if V_BOX_OPTIONS.from_stage:
if V_BOX_OPTIONS.from_stage == 'start':
from_stg_index = 0
else:
from_stg_index = stages.index(V_BOX_OPTIONS.from_stage)
if V_BOX_OPTIONS.to_stage:
if V_BOX_OPTIONS.from_stage == 'end':
to_stg_index = -1
else:
to_stg_index = stages.index(V_BOX_OPTIONS.to_stage) + 1
if to_stg_index is not None:
install_stages = stages[from_stg_index:to_stg_index]
else:
install_stages = stages[from_stg_index:]
LOG.info("Executing %s stage(s): %s.", len(install_stages), install_stages)
validate(V_BOX_OPTIONS, install_stages)
validate(V_BOX_OPTIONS, install_stages)
stg_no = 0
prev_stage = None
for stage in install_stages:
stg_no += 1
stg_start_time = time.time()
try:
stg_msg = f"({stg_no}/{len(install_stages)}) Entering stage {stage}"
log_heading_msg(stg_msg)
STAGE_CALLBACKS[stage][CALLBACK]()
stg_no = 0
prev_stage = None
for stage in install_stages:
stg_no += 1
stg_start_time = time.time()
try:
stg_msg = f"({stg_no}/{len(install_stages)}) Entering stage {stage}"
log_heading_msg(stg_msg)
STAGE_CALLBACKS[stage][CALLBACK]()
# Take snapshot if configured
if V_BOX_OPTIONS.snapshot:
vboxmanage.take_snapshot(
V_BOX_OPTIONS.labname,
f"snapshot-AFTER-{stage}")
# Take snapshot if configured
if V_BOX_OPTIONS.snapshot:
vboxmanage.take_snapshot(
V_BOX_OPTIONS.labname,
f"snapshot-AFTER-{stage}")
# Compute KPIs
stg_duration = time.time() - stg_start_time
kpi.set_kpi_metric(stage, stg_duration)
kpi.print_kpi(stage)
kpi.print_kpi('total')
except Exception as stg_exc:
stg_duration = time.time() - stg_start_time
kpi.set_kpi_metric(stage, stg_duration)
LOG.error("INSTALL FAILED, ABORTING!")
kpi.print_kpi_metrics()
LOG.info("Exception details: %s", repr(stg_exc))
raise
# Stage completed
prev_stage = stage
# Compute KPIs
stg_duration = time.time() - stg_start_time
kpi.set_kpi_metric(stage, stg_duration)
kpi.print_kpi(stage)
kpi.print_kpi('total')
except Exception as stg_exc:
stg_duration = time.time() - stg_start_time
kpi.set_kpi_metric(stage, stg_duration)
LOG.error("INSTALL FAILED, ABORTING!")
kpi.print_kpi_metrics()
LOG.info("Exception details: %s", repr(stg_exc))
raise
# Stage completed
prev_stage = stage
LOG.info("INSTALL SUCCEEDED!")
kpi.print_kpi_metrics()
LOG.info("INSTALL SUCCEEDED!")
kpi.print_kpi_metrics()
finally:
close_ssh_connections(SSH_CONNECTIONS)