Rework subprocess utility function

[deliverable/lttng-ivc.git] / lttng_ivc / utils / runtime.py

diff --git a/lttng_ivc/utils/runtime.py b/lttng_ivc/utils/runtime.py
index c2d0b8397299b5f0453bcb2550873e108aee28ea..a25d0bb97fbc77d711e43439fd7039bc1664baed 100644 (file)
--- a/lttng_ivc/utils/runtime.py
+++ b/lttng_ivc/utils/runtime.py
@@ -7,7 +7,6 @@ import logging
 import shutil
 import contextlib
 import pprint
-import traceback
 
 from tempfile import TemporaryDirectory
 
@@ -47,9 +46,10 @@ class Runtime(object):
         self._runtime_log_aggregation = os.path.join(self.__runtime_log, "runtime.log")
 
         self._run_command_count = 0
+        self._is_test_modules_loaded = False
 
         self.special_env_variables = {"LTTNG_UST_DEBUG": "1",
-                                      #"LTTNG_APP_SOCKET_TIMEOUT": "-1",
+                                      "LTTNG_APP_SOCKET_TIMEOUT": "-1",
                                       #"LTTNG_UST_REGISTER_TIMEOUT": "-1",
                                       "LTTNG_NETWORK_SOCKET_TIMEOUT": "-1"}
 
@@ -67,16 +67,26 @@ class Runtime(object):
     def add_project(self, project):
         self.__projects.append(project)
 
+    def remove_project(self, project):
+        self.__projects.remove(project)
+
     def subprocess_signal(self, subprocess_uuid, signal):
         self.__subproces[subprocess_uuid].send_signal(signal)
 
-    def subprocess_terminate(self, subprocess_uuid, timeout=60):
+    def subprocess_terminate(self, subprocess_uuid, timeout=60, check_return=True):
         process = self.__subprocess[subprocess_uuid]
         process.terminate()
-        process.wait(timeout)
+        try:
+            process.wait(timeout)
+        except subprocess.TimeoutExpired:
+            # Force kill
+            return self.subprocess_kill(process)
         stdout, stderr = self.__stdout_stderr[subprocess_uuid]
         stdout.close()
         stderr.close()
+        if check_return:
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(process.returncode, process.args)
         return process
 
     def subprocess_kill(self, subprocess_uuid):
@@ -88,6 +98,17 @@ class Runtime(object):
         stderr.close()
         return process
 
+    def subprocess_wait(self, subprocess_uuid, check_return=True):
+        process = self.__subprocess[subprocess_uuid]
+        process.wait()
+        stdout, stderr = self.__stdout_stderr[subprocess_uuid]
+        stdout.close()
+        stderr.close()
+        if check_return:
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(process.returncode, process.args)
+        return process
+
     def get_subprocess_stdout_path(self, subprocess_uuid):
         stdout, stderr = self.__stdout_stderr[subprocess_uuid]
         return stdout.name
@@ -114,13 +135,13 @@ class Runtime(object):
         with open(env_path, 'w') as env_out:
             pprint.pprint(env, stream=env_out)
 
-        p = subprocess.Popen(args, stdout=stdout, stderr=stderr, env=env)
+        p = subprocess.Popen(args, stdout=stdout, stderr=stderr, env=env, cwd=cwd)
         self.__subprocess[tmp_id] = p
         self.__stdout_stderr[tmp_id] = (stdout, stderr)
         _logger.debug("Spawned sub pid: {} args: {} stdout: {} stderr{}".format(p.pid, p.args, out_path, err_path))
         return tmp_id
 
-    def run(self, command_line, cwd=None, check_return=True):
+    def run(self, command_line, cwd=None, check_return=True, ld_preload="", classpath=""):
         """
         Run the command and return a tuple of a (CompletedProcess, stdout_path,
         stderr_path). The subprocess is already executed and returned. The
@@ -129,6 +150,12 @@ class Runtime(object):
         args = shlex.split(command_line)
         env = self.get_env()
 
+        if ld_preload:
+            env['LD_PRELOAD'] = ld_preload
+        if classpath:
+            env['CLASSPATH'] = classpath
+
+
         tmp_id = self._run_command_count
         self._run_command_count += 1
 
@@ -143,7 +170,8 @@ class Runtime(object):
 
         env_path = os.path.join(self.__runtime_log, str(tmp_id) + ".env")
         with open(env_path, 'w') as env_out:
-            pprint.pprint(env, stream=env_out)
+            for key, value in env.items():
+                env_out.write('{}={}\n'.format(key, value))
 
         cp = subprocess.run(args, stdout=stdout, stderr=stderr, env=env, cwd=cwd)
         _logger.debug("Command #{} args: {} stdout: {} stderr{}".format(tmp_id, cp.args, out_path, err_path))
@@ -183,7 +211,7 @@ class Runtime(object):
         library_path = []
         for project in self.__projects:
             library_path.append(project.get_ld_library_path())
-        return " ".join(library_path)
+        return ":".join(library_path)
 
     def get_bin_path(self):
         path = []
@@ -227,6 +255,16 @@ class Runtime(object):
                     env[var] = value
         return env
 
+    def load_test_module(self):
+        # Base directory is provided by env
+        self.run("modprobe lttng-test")
+        self._is_test_modules_loaded = True
+
+    def unload_test_module(self, check_return=True):
+        # Base directory is provided by env
+        if self._is_test_modules_loaded:
+            self.run("modprobe -r lttng-test", check_return=check_return)
+
     def close(self):
         for key, subp in self.__subprocess.items():
             subp.terminate()
@@ -237,7 +275,10 @@ class Runtime(object):
             stdout.close()
             stderr.close()
 
+        # Always try to remove test module but do not perform check on return
+        # value.
+        self.unload_test_module(False)
+
         # Copy the lttng_home used at runtime using hardlink to prevent useless
         # data duplication
         shutil.copytree(self.lttng_home, self.__post_runtime_lttng_home_path, copy_function=os.link)
-