Add tool_pid argument to parent_setup_fn

benchexec/baseexecutor.py

@@ -75,8 +75,8 @@ def _start_execution(
         parent_cleanup_fn,
     ):
         """Actually start the tool and the measurements.
-        @param parent_setup_fn a function without parameters that is called in the parent process
-            immediately before the tool is started
+        @param parent_setup_fn a function that is called in the parent process
+            immediately before the tool is started. The keyword-arguments passed may differ in subclasses.
         @param child_setup_fn a function without parameters that is called in the child process
             before the tool is started
         @param parent_cleanup_fn a function that is called in the parent process
@@ -88,14 +88,21 @@ def _start_execution(
             and the result of parent_cleanup_fn (do not use os.wait)
         """
 
+        from_child, to_parent = os.pipe()
+
         def pre_subprocess():
+            os.close(from_child)
+
             # Do some other setup the caller wants.
             child_setup_fn()
 
             # put us into the cgroup(s)
             pid = os.getpid()
             cgroups.add_task(pid)
 
+            os.write(to_parent, str(pid).encode())
+            os.close(to_parent)
+
         # Set HOME and TMPDIR to fresh directories.
         tmp_dir = os.path.join(temp_dir, "tmp")
         home_dir = os.path.join(temp_dir, "home")
@@ -108,8 +115,6 @@ def pre_subprocess():
         env["TEMP"] = tmp_dir
         logging.debug("Executing run with $HOME and $TMPDIR below %s.", temp_dir)
 
-        parent_setup = parent_setup_fn()
-
         p = subprocess.Popen(
             args,
             stdin=stdin,
@@ -118,8 +123,17 @@ def pre_subprocess():
             env=env,
             cwd=cwd,
             close_fds=True,
+            pass_fds=(to_parent,),
             preexec_fn=pre_subprocess,
         )
+        print("Continuing")
+
+        os.close(to_parent)
+
+        # read at most 10 bytes because this is enough for 32bit int
+        child_pid = int(os.read(from_child, 10))
+        os.close(from_child)
+        parent_setup = parent_setup_fn(child_pid=child_pid)
 
         def wait_and_get_result():
             exitcode, ru_child = self._wait_for_process(p.pid, args[0])

benchexec/containerexecutor.py

@@ -976,7 +976,10 @@ def check_child_exit_code():
 
             # start measurements
             cgroups.add_task(grandchild_pid)
-            parent_setup = parent_setup_fn()
+            parent_setup = parent_setup_fn(
+                child_pid=child_pid,
+                grandchild_pid=grandchild_pid,
+            )
 
             # Signal grandchild that setup is finished
             os.write(to_grandchild, MARKER_PARENT_COMPLETED)

benchexec/runexecutor.py

@@ -620,6 +620,8 @@ def execute_run(
         files_size_limit=None,
         error_filename=None,
         write_header=True,
+        parent_setup_fn=None,
+        parent_cleanup_fn=None,
         **kwargs,
     ) -> Dict[str, Any]:  # pytype: disable=signature-mismatch
         """
@@ -648,6 +650,8 @@ def execute_run(
         @param files_size_limit: None or maximum size of files that may be written.
         @param error_filename: the file where the error output should be written to (default: same as output_filename)
         @param write_headers: Write informational headers to the output and the error file if separate (default: True)
+        @param parent_setup_fn: A function that gets called in the parent process before the tool is allowed to run (default: noop fn). When used with use_namespaces=True, the function recieves child_pid and grandchild_pid.
+        @param parent_cleanup_fn: A function that gets called after the tool finishes running with the return value of parent_setup_fn.
         @param **kwargs: further arguments for ContainerExecutor.execute_run()
         @return: dict with result of run (measurement results and process exitcode)
         """
@@ -745,6 +749,13 @@ def execute_run(
             if files_size_limit < 0:
                 sys.exit(f"Invalid files-size limit {files_size_limit}.")
 
+        self.parent_setup_fn = (
+            util.dummy_fn if parent_setup_fn is None else parent_setup_fn
+        )
+        self.parent_cleanup_fn = (
+            util.dummy_fn if parent_cleanup_fn is None else parent_cleanup_fn
+        )
+
         try:
             return self._execute(
                 args,
@@ -829,19 +840,20 @@ def _execute(
                     # Disable energy measurements because we use only parts of a CPU
                     packages = None
 
-        def preParent():
+        def preParent(**kwargs):
             """Setup that is executed in the parent process immediately before the actual tool is started."""
             # start measurements
+            parent_setup = self.parent_setup_fn(**kwargs)
             if self._energy_measurement is not None and packages:
                 self._energy_measurement.start()
             starttime = util.read_local_time()
             walltime_before = time.monotonic()
-            return starttime, walltime_before
+            return starttime, walltime_before, parent_setup
 
         def postParent(preParent_result, exit_code, base_path):
             """Cleanup that is executed in the parent process immediately after the actual tool terminated."""
             # finish measurements
-            starttime, walltime_before = preParent_result
+            starttime, walltime_before, parent_setup = preParent_result
             walltime = time.monotonic() - walltime_before
             energy = (
                 self._energy_measurement.stop() if self._energy_measurement else None
@@ -869,6 +881,8 @@ def postParent(preParent_result, exit_code, base_path):
             if exit_code.value not in [0, 1]:
                 _get_debug_output_after_crash(output_filename, base_path)
 
+            self.parent_cleanup_fn(parent_setup, exit_code, base_path)
+
             return starttime, walltime, energy
 
         def preSubprocess():

benchexec/test_runexecutor.py

@@ -887,6 +887,35 @@ def test_frozen_process(self):
             "run output misses command output and was not executed properly",
         )
 
+    def test_parent_fns(self):
+        if not os.path.exists("/bin/sh"):
+            self.skipTest("missing /bin/sh")
+        parent_setup_ran = False
+        parent_cleanup_ran = False
+
+        def parent_setup_fn(*, child_pid, **kwargs):
+            # I don't want to require psutil just for this
+            # I'll just read the procfs
+            assert os.path.exists(f"/proc/{child_pid}")
+            nonlocal parent_setup_ran
+            parent_setup_ran = True
+            return 12345
+
+        def parent_cleanup_fn(parent_setup, exit_code, path):
+            assert parent_setup == 12345
+            nonlocal parent_cleanup_ran
+            parent_cleanup_ran = True
+
+        self.execute_run(
+            "/bin/sh",
+            "-c",
+            "echo hi",
+            parent_setup_fn=parent_setup_fn,
+            parent_cleanup_fn=parent_cleanup_fn,
+        )
+        assert parent_setup_ran
+        assert parent_cleanup_ran
+
 
 class TestRunExecutorWithContainer(TestRunExecutor):
     def setUp(self, *args, **kwargs):
@@ -1092,6 +1121,36 @@ def test_result_file_log_limit(self):
         count_msg = next(msg for msg in log.output if " output files matched" in msg)
         self.assertIn(f"{file_count} output files matched", count_msg)
 
+    def test_parent_fns(self):
+        if not os.path.exists("/bin/sh"):
+            self.skipTest("missing /bin/sh")
+        parent_setup_ran = False
+        parent_cleanup_ran = False
+
+        def parent_setup_fn(*, grandchild_pid, child_pid, **kwargs):
+            # I don't want to require psutil just for this
+            # I'll just read the procfs
+            assert os.path.exists(f"/proc/{grandchild_pid}")
+            assert os.path.exists(f"/proc/{child_pid}")
+            nonlocal parent_setup_ran
+            parent_setup_ran = True
+            return 12345
+
+        def parent_cleanup_fn(parent_setup, exit_code, path):
+            assert parent_setup == 12345
+            nonlocal parent_cleanup_ran
+            parent_cleanup_ran = True
+
+        self.execute_run(
+            "/bin/sh",
+            "-c",
+            "echo hi",
+            parent_setup_fn=parent_setup_fn,
+            parent_cleanup_fn=parent_cleanup_fn,
+        )
+        assert parent_setup_ran
+        assert parent_cleanup_ran
+
     def test_file_count_limit(self):
         if not os.path.exists("/bin/sh"):
             self.skipTest("missing /bin/sh")