diff --git a/scripts/rpc.py b/scripts/rpc.py index f328e8a03..1dfd263b6 100755 --- a/scripts/rpc.py +++ b/scripts/rpc.py @@ -1737,10 +1737,11 @@ Format: 'user:u1 secret:s1 muser:mu1 msecret:ms1,user:u2 secret:s2 muser:mu2 mse args = parser.parse_args() - try: - args.client = rpc.client.JSONRPCClient(args.server_addr, args.port, args.timeout, log_level=getattr(logging, args.verbose.upper())) - args.func(args) - except JSONRPCException as ex: - print("Exception:") - print(ex.message) - exit(1) + with rpc.client.JSONRPCClient(args.server_addr, args.port, args.timeout, log_level=getattr(logging, args.verbose.upper())) as client: + try: + args.client = client + args.func(args) + except JSONRPCException as ex: + print("Exception:") + print(ex.message) + exit(1) diff --git a/scripts/rpc/client.py b/scripts/rpc/client.py index 967ea7849..e948fc830 100644 --- a/scripts/rpc/client.py +++ b/scripts/rpc/client.py @@ -48,6 +48,12 @@ class JSONRPCClient(object): raise JSONRPCException("Error while connecting to %s\n" "Error details: %s" % (addr, ex)) + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, traceback): + self.close() + def get_logger(self): return self._logger @@ -61,9 +67,11 @@ class JSONRPCClient(object): self._logger.setLevel(lvl) self._logger.info("Log level set to %s", lvl) - def __del__(self): + def close(self): if getattr(self, "sock", None): + self.sock.shutdown(socket.SHUT_RDWR) self.sock.close() + self.sock = None def add_request(self, method, params): self._request_id += 1 diff --git a/scripts/spdkcli.py b/scripts/spdkcli.py index 7dcf76bc6..743ed0607 100755 --- a/scripts/spdkcli.py +++ b/scripts/spdkcli.py @@ -6,6 +6,7 @@ from os import getuid from rpc.client import JSONRPCException from configshell_fb import ConfigShell, shell, ExecutionError from spdkcli import UIRoot +import rpc.client from pyparsing import (alphanums, Optional, Suppress, Word, Regex, removeQuotes, dblQuotedString, OneOrMore) @@ -43,29 +44,30 @@ def main(): help="commands to execute by SPDKCli as one-line command") args = parser.parse_args() - root_node = UIRoot(args.socket, spdk_shell) - root_node.verbose = args.verbose - try: - root_node.refresh() - except BaseException: - pass - - if len(args.commands) > 0: + with rpc.client.JSONRPCClient(args.socket) as client: + root_node = UIRoot(client, spdk_shell) + root_node.verbose = args.verbose try: - spdk_shell.interactive = False - spdk_shell.run_cmdline(" ".join(args.commands)) - except Exception as e: - sys.stderr.write("%s\n" % e) - sys.exit(1) - sys.exit(0) + root_node.refresh() + except BaseException: + pass - spdk_shell.con.display("SPDK CLI v0.1") - spdk_shell.con.display("") - while not spdk_shell._exit: - try: - spdk_shell.run_interactive() - except (JSONRPCException, ExecutionError) as e: - spdk_shell.log.error("%s" % e) + if len(args.commands) > 0: + try: + spdk_shell.interactive = False + spdk_shell.run_cmdline(" ".join(args.commands)) + except Exception as e: + sys.stderr.write("%s\n" % e) + sys.exit(1) + sys.exit(0) + + spdk_shell.con.display("SPDK CLI v0.1") + spdk_shell.con.display("") + while not spdk_shell._exit: + try: + spdk_shell.run_interactive() + except (JSONRPCException, ExecutionError) as e: + spdk_shell.log.error("%s" % e) if __name__ == "__main__": diff --git a/scripts/spdkcli/ui_root.py b/scripts/spdkcli/ui_root.py index 38ad48c9a..3c19adc95 100644 --- a/scripts/spdkcli/ui_root.py +++ b/scripts/spdkcli/ui_root.py @@ -10,14 +10,14 @@ class UIRoot(UINode): """ Root node for CLI menu tree structure. Refreshes running config on startup. """ - def __init__(self, s, shell): + def __init__(self, client, shell): UINode.__init__(self, "/", shell=shell) self.current_bdevs = [] self.current_lvol_stores = [] self.current_vhost_ctrls = [] self.current_nvmf_transports = [] self.current_nvmf_subsystems = [] - self.set_rpc_target(s) + self.set_rpc_target(client) self.verbose = False self.is_init = self.check_init() self.methods = [] @@ -45,8 +45,8 @@ class UIRoot(UINode): if self.has_subsystem("iscsi"): UIISCSI(self) - def set_rpc_target(self, s): - self.client = rpc.client.JSONRPCClient(s) + def set_rpc_target(self, client): + self.client = client def print_array(self, a): return " ".join(a) diff --git a/test/json_config/clear_config.py b/test/json_config/clear_config.py index decf3eae3..e86f41b90 100755 --- a/test/json_config/clear_config.py +++ b/test/json_config/clear_config.py @@ -209,9 +209,10 @@ if __name__ == "__main__": args = parser.parse_args() - try: - args.client = rpc.client.JSONRPCClient(args.server_addr, args.port, args.timeout, log_level=getattr(logging, args.verbose.upper())) - except JSONRPCException as ex: - print((ex.message)) - exit(1) - args.func(args) + with rpc.client.JSONRPCClient(args.server_addr, args.port, args.timeout, log_level=getattr(logging, args.verbose.upper())) as client: + try: + args.client = client + args.func(args) + except JSONRPCException as ex: + print((ex.message)) + exit(1)