From 534a69a1797944d2a6ce195e66cbca441fd5d656 Mon Sep 17 00:00:00 2001
From: Pierre Lasorak <plasorak@fnal.gov>
Date: Fri, 17 Jan 2025 16:01:46 +0100
Subject: [PATCH] add a wrapper to the processes executed over ssh

---
 setup.cfg                                     |  3 +-
 .../apps/__main_ssh_process_wrapper__.py      | 33 +++++++++++++++++++
 .../process_manager/ssh_process_manager.py    |  2 +-
 3 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 src/drunc/apps/__main_ssh_process_wrapper__.py

diff --git a/setup.cfg b/setup.cfg
index e6e6f8e2..e6e997e0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -30,4 +30,5 @@ console_scripts =
     drunc-unified-shell = drunc.apps.__main_unified_shell__:main
     drunc-fsm-tests = drunc.tests.__main_fsm__:main
     application-registry-service = drunc.apps.__main_app_connectivity_server__:main
-    drunc-ssh-validator = drunc.apps.__main_ssh_validator__:main
\ No newline at end of file
+    drunc-ssh-validator = drunc.apps.__main_ssh_validator__:main
+    drunc-ssh-process-wrapper = drunc.apps.__main_ssh_process_wrapper__:main
\ No newline at end of file
diff --git a/src/drunc/apps/__main_ssh_process_wrapper__.py b/src/drunc/apps/__main_ssh_process_wrapper__.py
new file mode 100644
index 00000000..2eda4c40
--- /dev/null
+++ b/src/drunc/apps/__main_ssh_process_wrapper__.py
@@ -0,0 +1,33 @@
+import signal
+import os
+import time
+import subprocess
+import click
+
+def terminate_all(sig, frame):
+    print('')
+    pgrp = os.getpgid(os.getpid())
+    os.killpg(pgrp, signal.SIGKILL)
+
+
+@click.command()
+@click.argument("cmd")
+@click.option("-l", "--log", "log_path", type=click.Path(file_okay=True, dir_okay=False))
+def launch(cmd: str, log_path: str):
+    signal.signal(signal.SIGHUP, terminate_all)
+
+    with open(log_path, 'w') as logfile:
+        proc = subprocess.Popen(cmd, shell=True, stdout=logfile, stderr=logfile)
+        print(f"Started process with pid {proc.pid}")
+
+        return_code = None
+        while True:
+            return_code = proc.poll()
+            if return_code is not None:
+                break
+            time.sleep(0.1)
+
+
+
+if __name__ == '__main__':
+    launch()
\ No newline at end of file
diff --git a/src/drunc/process_manager/ssh_process_manager.py b/src/drunc/process_manager/ssh_process_manager.py
index 5f0d844c..37e1653a 100644
--- a/src/drunc/process_manager/ssh_process_manager.py
+++ b/src/drunc/process_manager/ssh_process_manager.py
@@ -268,7 +268,7 @@ def __boot(self, boot_request:BootRequest, uuid:str) -> ProcessInstance:
                 if cmd[-1] == ';':
                     cmd = cmd[:-1]
 
-                arguments = [user_host, "-tt", "-o StrictHostKeyChecking=no", f'{{ {cmd} ; }} &> {log_file}']
+                arguments = [user_host, "-tt", "-o StrictHostKeyChecking=no", f'{cmd_env} drunc-ssh-process-wrapper --log {log_file} {cmd}']
                 self.log.debug(f"{arguments}")
                 # arguments = [user_host, "-tt", "-o StrictHostKeyChecking=no", f'{{ {cmd} ; }} > >(tee -a {log_file}) 2> >(tee -a {log_file} >&2)']
                 # I'm gonna bail now and read that log file, anyway, it's probably better that heavy logger applications don't clog up the process manager CPU.