From a8cacf0c3ff83324b38fd425edf30e37da0da9bf Mon Sep 17 00:00:00 2001
From: Luke Logan <lukemartinlogan@gmail.com>
Date: Wed, 29 Jan 2025 03:58:32 -0600
Subject: [PATCH] Support LD_PRELOAD with gdbserver

---
 jarvis_util/shell/exec_info.py  | 17 +++++++++++++----
 jarvis_util/shell/local_exec.py | 18 ++++++++++--------
 jarvis_util/shell/mpi_exec.py   |  2 +-
 jarvis_util/shell/pssh_exec.py  |  2 +-
 4 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/jarvis_util/shell/exec_info.py b/jarvis_util/shell/exec_info.py
index 676bff9..bd8cf71 100644
--- a/jarvis_util/shell/exec_info.py
+++ b/jarvis_util/shell/exec_info.py
@@ -241,13 +241,22 @@ def set_exit_code_list(self, nodes):
             if node.exit_code:
                 self.exit_code = node.exit_code
 
-    def get_dbg_cmd(self, cmd, dbg_port):
+    def get_dbg_cmd(self, cmd, exec_info):
         """
         Get the command to debug a program
 
         :param cmd: the command to debug
-        :param dbg_port: the debugging port to use
-        :return:
+        :param exec_info: exec information
+        :return: the debug command
         """
-        return f'gdbserver localhost:{dbg_port} {cmd}'
+        dbg_port = exec_info.dbg_port
+        preload = ""
+        if 'LD_PRELOAD' in exec_info.env:
+            exec_info.env = exec_info.env.copy()
+            preload = exec_info.env['LD_PRELOAD']
+            del exec_info.env['LD_PRELOAD']
+        if len(preload):
+            return f'gdbserver localhost:{dbg_port} env LD_PRELOAD={preload} {cmd}'
+        else:
+            return f'gdbserver localhost:{dbg_port} {cmd}'
 
diff --git a/jarvis_util/shell/local_exec.py b/jarvis_util/shell/local_exec.py
index a4c27f7..839f712 100644
--- a/jarvis_util/shell/local_exec.py
+++ b/jarvis_util/shell/local_exec.py
@@ -53,13 +53,6 @@ def __init__(self, cmd, exec_info):
         self.print_stderr_thread = None
         self.exit_code = 0
 
-        # Copy ENV
-        self.basic_env = exec_info.basic_env.copy()
-        self.env = exec_info.env.copy()
-        for key, val in os.environ.items():
-            if key not in self.env:
-                self.env[key] = val
-
         # Managing command execution
         self.sudo = exec_info.sudo
         self.stdin = exec_info.stdin
@@ -69,12 +62,21 @@ def __init__(self, cmd, exec_info):
             self.cwd = os.getcwd()
         else:
             self.cwd = exec_info.cwd
+        self.basic_env = exec_info.basic_env.copy()
 
         # Create the command
         cmd = self.smash_cmd(cmd, self.sudo, self.basic_env, exec_info.sudoenv)
         if exec_info.do_dbg:
-            cmd = self.get_dbg_cmd(cmd, exec_info.dbg_port)
+            cmd = self.get_dbg_cmd(cmd, exec_info)
         self.cmd = cmd
+
+        # Copy ENV
+        self.env = exec_info.env.copy()
+        for key, val in os.environ.items():
+            if key not in self.env:
+                self.env[key] = val
+
+        # Execute the command
         if self.jutil.debug_local_exec:
             print(cmd)
         self._start_bash_processes()
diff --git a/jarvis_util/shell/mpi_exec.py b/jarvis_util/shell/mpi_exec.py
index 03473ee..a13ac4d 100644
--- a/jarvis_util/shell/mpi_exec.py
+++ b/jarvis_util/shell/mpi_exec.py
@@ -48,7 +48,7 @@ def __init__(self, cmd, exec_info):
         self.mpi_env = exec_info.env
         if exec_info.do_dbg:
             self.base_cmd = cmd # To append to the extra processes
-            self.cmd = self.get_dbg_cmd(cmd, exec_info.dbg_port)
+            self.cmd = self.get_dbg_cmd(cmd, exec_info)
         super().__init__(self.mpicmd(),
                          exec_info.mod(env=exec_info.basic_env,
                                        do_dbg=False))
diff --git a/jarvis_util/shell/pssh_exec.py b/jarvis_util/shell/pssh_exec.py
index 7fd7cf8..70da4ef 100644
--- a/jarvis_util/shell/pssh_exec.py
+++ b/jarvis_util/shell/pssh_exec.py
@@ -31,7 +31,7 @@ def __init__(self, cmd, exec_info):
         if not self.is_local:
             dbg_cmd = cmd
             if exec_info.do_dbg:
-                dbg_cmd = self.get_dbg_cmd(cmd, exec_info.dbg_port)
+                dbg_cmd = self.get_dbg_cmd(cmd, exec_info)
             for i, host in enumerate(self.hosts):
                 sshcmd = cmd
                 if i == 0: