Pawel plesniak emuhammad/split shell fixing

src/drunc/controller/controller.py

@@ -1,4 +1,4 @@
-import multiprocessing
+import os
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
@@ -129,17 +129,29 @@ def __init__(self, configuration, name: str, session: str, token: Token):
         self.connectivity_service_thread = None
         self.uri = ""
         if self.configuration.session.connectivity_service:
-            connection_server = self.configuration.session.connectivity_service.host
-            connection_port = (
-                self.configuration.session.connectivity_service.service.port
+            # Remaps the localhost into the correct server
+            # and also grabs the correct port from the right environment from the config
+
+            connection_server_host = (
+                self.configuration.session.connectivity_service.host
             )
+            connection_port = os.getenv("CONNECTION_PORT")
+            if connection_server_host == "localhost":
+                injected_hostname = os.getenv("DRUNC_HOST_NAME")
+                if not injected_hostname:
+                    raise ValueError("DRUNC_HOST_NAME environment variable is not set.")
+                self.log.debug(
+                    f"Remapping connectivity service host from 'localhost' to '{injected_hostname}'"
+                )
+                connection_server_host = injected_hostname
+
             log_init.info(
-                f"Connectivity server {connection_server}:{connection_port} is enabled"
+                f"Connectivity server {connection_server_host}:{connection_port} is enabled"
             )
 
             self.connectivity_service = ConnectivityServiceClient(
                 session=self.session,
-                address=f"{connection_server}:{connection_port}",
+                address=f"{connection_server_host}:{connection_port}",
             )
 
     def init_controller(self) -> None:
@@ -326,6 +338,13 @@ def advertise_control_address(self, address):
         if not self.connectivity_service:
             return
 
+        self.log.debug(
+            f"Looking for connectivity service at address {self.connectivity_service.address}"
+        )
+        if not self.connectivity_service.is_ready(timeout=20):
+            raise ValueError(
+                "Connectivity service unavailable for control address advertising."
+            )
         self.log.info(
             f"Registering {self.name} ({address}) to the connectivity service at {self.connectivity_service.address}"
         )
@@ -384,14 +403,6 @@ def terminate(self):
             except Exception as e:
                 self.log.warning(f"Error stopping opmon publisher: {e}")
 
-        self.log.debug("Threading threads")
-        for t in threading.enumerate():
-            self.log.debug(f"{t.name} TID: {t.native_id} is_alive: {t.is_alive}")
-
-        with multiprocessing.Manager() as manager:
-            self.log.debug("Multiprocess threads")
-            self.log.debug(manager.list())
-
     def __del__(self):
         self.terminate()
 

src/drunc/controller/interface/context.py

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

src/drunc/process_manager/configuration.py

@@ -15,7 +15,7 @@
 from drunc.exceptions import DruncCommandException
 from drunc.process_manager.exceptions import UnknownProcessManagerType
 from drunc.utils.configuration import ConfHandler
-from drunc.utils.utils import get_logger
+from drunc.utils.utils import get_logger, touch_and_chmod
 
 if TYPE_CHECKING:
     import conffwk
@@ -110,6 +110,7 @@ def _post_process_oks(self) -> None:
                     opmon_conf,
                 )
             else:
+                touch_and_chmod(opmon_conf.path)
                 self.opmon_publisher = OpMonPublisher(
                     conf=opmon_conf, rich_handler=True
                 )

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,8 +12,8 @@
 )
 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.utils import get_logger
+from drunc.utils.shell_utils import InterruptedCommand, log_pm_cmd
+from drunc.utils.utils import get_logger, resolve_context_peer
 
 
 @click.command("boot")
@@ -43,11 +44,15 @@ def boot(
     override_logs: bool,
 ) -> None:
     log = get_logger("process_manager.shell")
-    processes = obj.get_driver("process_manager").ps(ProcessQuery(user=user))
+    log_pm_cmd(obj)
+    processes = obj.get_driver("process_manager").ps(
+        ProcessQuery(user=user, session=session_name)
+    )
 
+    # The run control will validate this in the session manager in the future
     if len(processes.values) > 0:
         click.confirm(
-            f"You already have {len(processes.values)} processes running, are you sure you want to boot a session?",
+            f"You already have {len(processes.values)} processes running for {session_name}, are you sure you want to boot a session?",
             abort=True,
         )
 
@@ -76,6 +81,7 @@ def boot(
         raise e
 
     controller_address = obj.get_driver("process_manager").controller_address
+    controller_address = resolve_context_peer(controller_address)
     if controller_address:
         obj.print(
             Panel(
@@ -129,6 +135,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 +157,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 +178,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 +189,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 +228,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,18 +263,28 @@ 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")
@@ -276,6 +326,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 +361,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

@@ -10,7 +10,9 @@
 from drunc.utils.utils import resolve_localhost_to_hostname
 
 
-class ProcessManagerContext(ShellContext):
+class ProcessManagerContext(ShellContext):  # boilerplatefest
+    shell_id = "process_manager_shell"
+
     def __init__(self, *args, **kwargs):
         self.status_receiver = None
         super(ProcessManagerContext, self).__init__(*args, **kwargs)