From 600f955fb97e5b75f0c46a1655d1670f8f402c2f Mon Sep 17 00:00:00 2001 From: Wanyun Su Date: Wed, 18 Dec 2024 13:11:25 +0100 Subject: [PATCH] delete controller when terminate --- src/drunc/process_manager/interface/commands.py | 3 +++ src/drunc/utils/shell_utils.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/src/drunc/process_manager/interface/commands.py b/src/drunc/process_manager/interface/commands.py index 2ea6d90c..88f7a727 100644 --- a/src/drunc/process_manager/interface/commands.py +++ b/src/drunc/process_manager/interface/commands.py @@ -86,6 +86,8 @@ async def terminate(obj:ProcessManagerContext) -> None: from drunc.process_manager.utils import tabulate_process_instance_list obj.print(tabulate_process_instance_list(result.data, 'Terminated process', False)) + obj.delete_driver('controller') + @click.command('kill') @add_query_options(at_least_one=True) @click.pass_obj @@ -100,6 +102,7 @@ async def kill(obj:ProcessManagerContext, query:ProcessQuery) -> None: from drunc.process_manager.utils import tabulate_process_instance_list obj.print(tabulate_process_instance_list(result.data, 'Killed process', False)) + obj.delete_driver('controller') @click.command('flush') @add_query_options(at_least_one=False, all_processes_by_default=True) diff --git a/src/drunc/utils/shell_utils.py b/src/drunc/utils/shell_utils.py index 1b298c99..7f539c71 100644 --- a/src/drunc/utils/shell_utils.py +++ b/src/drunc/utils/shell_utils.py @@ -276,6 +276,14 @@ def get_driver(self, name:str=None) -> GRPCDriver: self._log.error(f'Controller-specific commands cannot be sent until the session is booted') raise SystemExit(1) # used to avoid having to catch multiple Attribute errors when this function gets called + def delete_driver(self, name: str) -> None: + if name in self._drivers: + del self._drivers[name] + self._log.info(f"Driver '{name}' has been deleted.") + else: + self._log.warn(f"Driver '{name}' does not exist in this context.") + + def get_token(self) -> Token: return self._token