Pawel plesniak emuhammad/split shell fixing

src/drunc/controller/interface/context.py

@@ -14,6 +14,8 @@
 
 
 class ControllerContext(ShellContext):  # boilerplatefest
+    shell_id = "controller_shell"
+
     def __init__(self):
         self.status_receiver = None
         self.took_control = False

src/drunc/process_manager/interface/cli_argument.py

@@ -7,15 +7,10 @@ def validate_conf_string(ctx, param, boot_configuration):
     return boot_configuration
 
 
-def add_query_options(at_least_one: bool, all_processes_by_default: bool = False):
-    def wrapper(f0):
-        f1 = click.option(
-            "-s",
-            "--session",
-            type=str,
-            default=None,
-            help="Select the processes on a particular session",
-        )(f0)
+def add_query_options_no_session(
+    at_least_one: bool, all_processes_by_default: bool = False
+):
+    def wrapper(f1):
         f2 = click.option(
             "-n",
             "--name",
@@ -41,3 +36,17 @@ def wrapper(f0):
         return generate_process_query(f4, at_least_one, all_processes_by_default)
 
     return wrapper
+
+
+def add_query_options(at_least_one: bool, all_processes_by_default: bool = False):
+    def wrapper(f0):
+        f1 = click.option(
+            "-s",
+            "--session",
+            type=str,
+            default=None,
+            help="Select the processes on a particular session",
+        )(f0)
+        return add_query_options_no_session(at_least_one, all_processes_by_default)(f1)
+
+    return wrapper

src/drunc/process_manager/interface/commands.py

@@ -1,4 +1,5 @@
 import getpass
+from time import sleep
 
 import click
 from druncschema.process_manager_pb2 import LogRequest, ProcessQuery
@@ -11,7 +12,7 @@
 )
 from drunc.process_manager.interface.context import ProcessManagerContext
 from drunc.process_manager.utils import tabulate_process_instance_list
-from drunc.utils.shell_utils import InterruptedCommand
+from drunc.utils.shell_utils import InterruptedCommand, log_pm_cmd
 from drunc.utils.utils import get_logger
 
 
@@ -43,6 +44,7 @@ def boot(
     override_logs: bool,
 ) -> None:
     log = get_logger("process_manager.shell")
+    log_pm_cmd(obj)
     processes = obj.get_driver("process_manager").ps(ProcessQuery(user=user))
 
     if len(processes.values) > 0:
@@ -129,6 +131,7 @@ def dummy_boot(
     session_name: str,
 ) -> None:
     log = get_logger("process_manager.shell")
+    log_pm_cmd(obj)
     log.debug(
         f"Running dummy_boot with {n_processes} processes for {sleep} seconds {n_sleeps} times, requested by user {user}"
     )
@@ -150,6 +153,16 @@ def dummy_boot(
         return
 
 
+@click.command("wait")
+@click.argument("sleep_time", type=int, default=1)
+@click.pass_obj
+def wait(obj: ProcessManagerContext, sleep_time: int) -> None:
+    log = get_logger("process_manager.wait")
+    log.info(f"Command [green]wait[/green] running for {sleep_time} seconds.")
+    sleep(sleep_time)  # seconds
+    log.info(f"Command [green]wait[/green] ran for {sleep_time} seconds.")
+
+
 @click.command("terminate")
 @click.option(
     "-w",
@@ -161,6 +174,7 @@ def dummy_boot(
 @click.pass_obj
 def terminate(obj: ProcessManagerContext, width: int | None) -> None:
     log = get_logger("process_manager.shell")
+    log_pm_cmd(obj)
     log.debug("Terminating")
     result = obj.get_driver("process_manager").terminate()
     if not result:
@@ -171,23 +185,35 @@ def terminate(obj: ProcessManagerContext, width: int | None) -> None:
     obj.delete_driver("controller")
 
 
+def kill_decorators(f):
+    f = click.pass_obj(f)
+    f = click.option(
+        "-w",
+        "--width",
+        type=int,
+        default=None,
+        help="Table width. Default is automatically calculated",
+    )(f)
+    f = click.option(
+        "--crash",
+        is_flag=True,
+        default=False,
+        help="Simulate a crash: send SIGKILL without any cleanup, leaving the process manager in an unexpected-death state.",
+    )(f)
+    return f
+
+
 @click.command("kill")
-@click.option(
-    "-w",
-    "--width",
-    type=int,
-    default=None,
-    help="Table width. Default is automatically calculated",
-)
 @add_query_options(at_least_one=True)
-@click.option(
-    "--crash",
-    is_flag=True,
-    default=False,
-    help="Simulate a crash: send SIGKILL without any cleanup, leaving the process manager in an unexpected-death state.",
-)
-@click.pass_obj
-def kill(obj: ProcessManagerContext, query: ProcessQuery, width: int | None) -> None:
+@kill_decorators
+def kill(obj, query, width):
+    log_pm_cmd(obj)
+    return kill_impl(obj, query, width)
+
+
+def kill_impl(
+    obj: ProcessManagerContext, query: ProcessQuery, width: int | None
+) -> None:
     log = get_logger("process_manager.shell")
     log.debug(f"Killing with query {query}")
     result = obj.get_driver("process_manager").kill(query)
@@ -198,17 +224,27 @@ def kill(obj: ProcessManagerContext, query: ProcessQuery, width: int | None) ->
     )  # rich tables require console printing
 
 
+def flush_decorators(f):
+    f = click.pass_obj(f)
+    f = click.option(
+        "-w",
+        "--width",
+        type=int,
+        default=None,
+        help="Table width. Default is automatically calculated",
+    )(f)
+    return f
+
+
 @click.command("flush")
-@click.option(
-    "-w",
-    "--width",
-    type=int,
-    default=None,
-    help="Table width. Default is automatically calculated",
-)
 @add_query_options(at_least_one=False, all_processes_by_default=True)
-@click.pass_obj
-def flush(
+@flush_decorators
+def flush(obj, query, width):
+    log_pm_cmd(obj)
+    return flush_impl(obj, query, width)
+
+
+def flush_impl(
     obj: ProcessManagerContext,
     query: ProcessQuery,
     width: int | None,
@@ -223,22 +259,34 @@ def flush(
     )  # rich tables require console printing
 
 
+def logs_decorators(f):
+    f = click.pass_obj(f)
+    f = click.option("--grep", type=str, default=None)(f)
+    f = click.option(
+        "--how-far",
+        type=int,
+        show_default=True,
+        default=100,
+        help="How many lines one wants",
+    )(f)
+    return f
+
+
 @click.command("logs")
 @add_query_options(at_least_one=True)
-@click.option(
-    "--how-far",
-    type=int,
-    show_default=True,
-    default=100,
-    help="How many lines one wants",
-)
-@click.option("--grep", type=str, default=None)
-@click.pass_obj
-def logs(
+@logs_decorators
+def logs(obj, how_far, grep, query):
+    log_pm_cmd(obj)
+    return logs_impl(obj, how_far, grep, query)
+
+
+def logs_impl(
     obj: ProcessManagerContext, how_far: int, grep: str, query: ProcessQuery
 ) -> None:
     log = get_logger("process_manager.shell")
-    log.debug(f"Running logs with query {query}")
+    # TODO: MOVE BACK TO DEBUG BEFORE MERGE
+    # THIS IS USEFUL FOR TESTING THOUGH
+    log.error(f"Running logs with query {query}")
     log_req = LogRequest(
         how_far=how_far,
         query=query,
@@ -276,6 +324,11 @@ def logs(
 @add_query_options(at_least_one=True)
 @click.pass_obj
 def restart(obj: ProcessManagerContext, query: ProcessQuery) -> None:
+    log_pm_cmd(obj)
+    return restart_impl(obj, query)
+
+
+def restart_impl(obj: ProcessManagerContext, query: ProcessQuery) -> None:
     log = get_logger("process_manager.shell")
     log.debug(f"Restarting with query {query}")
     obj.get_driver("process_manager").restart(query)
@@ -306,6 +359,7 @@ def ps(
     width: int | None,
 ) -> None:
     log = get_logger("process_manager.shell")
+    log_pm_cmd(obj)
     log.debug(f"Running ps with query {query}")
     results = obj.get_driver("process_manager").ps(query)
     if not results:

src/drunc/process_manager/interface/context.py

@@ -14,6 +14,8 @@
 
 
 class ProcessManagerContext(ShellContext):  # boilerplatefest
+    shell_id = "process_manager_shell"
+
     def __init__(self, *args, **kwargs):
         self.status_receiver = None
         super(ProcessManagerContext, self).__init__(*args, **kwargs)

src/drunc/process_manager/interface/shell.py

@@ -14,6 +14,7 @@
     ps,
     restart,
     terminate,
+    wait,
 )
 from drunc.utils.grpc_utils import ServerUnreachable
 from drunc.utils.utils import (
@@ -59,6 +60,10 @@ def process_manager_shell(ctx, process_manager_address: str, log_level: str) ->
         # process_manager_shell_log.error(e.message) # TODO: Keep this for production branch, remove this from dev branch
         exit(1)
 
+    ctx.obj.get_driver("process_manager").send_msg(
+        f"{getpass.getuser()} connected from {ctx.obj.shell_id}"
+    )
+
     # Manually add file handler to process manager log
     # Not possible to initialise logger immediately as it requires
     # knowledge of the log path
@@ -75,6 +80,9 @@ def process_manager_shell(ctx, process_manager_address: str, log_level: str) ->
         ctx.obj.start_listening(desc.broadcast)
 
     def cleanup():
+        ctx.obj.get_driver("process_manager").send_msg(
+            f"{getpass.getuser()} disconnected from {ctx.obj.shell_id}"
+        )
         ctx.obj.terminate()
         process_manager_log.warning(
             f"[green]{getpass.getuser()}[/green] disconnected from the process manager through a [green]drunc-process-manager-shell[/green]"
@@ -83,6 +91,7 @@ def cleanup():
     ctx.call_on_close(cleanup)
 
     ctx.command.add_command(boot, "boot")
+    ctx.command.add_command(wait, "wait")
     ctx.command.add_command(terminate, "terminate")
     ctx.command.add_command(kill, "kill")
     ctx.command.add_command(flush, "flush")

src/drunc/process_manager/k8s_process_manager.py

@@ -13,6 +13,7 @@
 
 # Local Application Imports
 from druncschema.broadcast_pb2 import BroadcastType
+from druncschema.generic_pb2 import OutcomeFlag, OutcomeStatus
 from druncschema.process_manager_pb2 import (
     BootRequest,
     LogLines,
@@ -39,6 +40,7 @@
 from drunc.process_manager.configuration import (
     PROCESS_SHUTDOWN_ORDERING,
     ProcessManagerConfHandler,
+    ProcessManagerTypes,
 )
 from drunc.process_manager.process_manager import ProcessManager
 from drunc.process_manager.utils import on_parent_exit, validate_k8s_session_name
@@ -152,6 +154,8 @@ def run(self) -> None:
 
 
 class K8sProcessManager(ProcessManager):
+    pm_type = ProcessManagerTypes.K8s
+
     def __init__(self, configuration: ProcessManagerConfHandler, **kwargs) -> None:
         """
         Manages processes as Kubernetes Pods.
@@ -1910,6 +1914,17 @@ def _logs_impl(self, log_request: LogRequest) -> LogLines:
                 lines=[f"Could not retrieve logs: {e.reason}"],
             )
 
+    def _send_msg_impl(self, msg: str, peer: str) -> OutcomeStatus:
+        # Note: currently exact same implementation as ssh manager
+        # Although there is room here to change as necessary
+        try:
+            self.log.info(f"{msg}; from {peer}")
+        except Exception as e:
+            self.log.critical(f"Failed to receive message with exception {e}")
+            return OutcomeStatus(flag=OutcomeFlag.FAIL)
+
+        return OutcomeStatus(flag=OutcomeFlag.SUCCESS)
+
     def _boot_impl(self, boot_request: BootRequest) -> ProcessInstanceList:
         """
         Handles the 'boot' command from the gRPC interface.
@@ -2516,7 +2531,9 @@ def kill_and_wait(uuids, grace_period=None) -> None:
                 pods_by_role[role].append(uuid)
                 if role == "segment-controller":
                     tree_id = pod.metadata.labels.get(tree_id_label_key, "")
-                    segment_controller_depths[uuid] = tree_id.count(".") if tree_id else 0
+                    segment_controller_depths[uuid] = (
+                        tree_id.count(".") if tree_id else 0
+                    )
 
         # Kill in stages using our sorted lists
         for role in PROCESS_SHUTDOWN_ORDERING:

src/drunc/process_manager/oks_parser.py

@@ -5,7 +5,7 @@
 
 from drunc.exceptions import DruncException, DruncSetupException
 from drunc.process_manager.configuration import get_commandline_parameters
-from drunc.utils.utils import get_logger
+from drunc.utils.utils import file_is_read_only, get_logger
 
 if TYPE_CHECKING:
     import conffwk
@@ -75,8 +75,13 @@ def get_full_db_path(db_path: str) -> str:
         err_str = f"No files found in DUNEDAQ_DB_PATH matching {db_path}."
         raise DruncSetupException(err_str)
 
-    # If multiple matches are found, take the first instance that matches.
+    # Prefer the first writable match; if every match is read-only, fall back to the first one.
     resolved_path = unique_matched_files[0]
+    for matched_file in unique_matched_files:
+        if not file_is_read_only(matched_file):
+            resolved_path = matched_file
+            break
+
     log.debug(f"Path {db_path} resolved to {resolved_path}")
     return resolved_path