Add tool_pid argument to parent_setup_fn

benchexec/baseexecutor.py

@@ -49,7 +49,24 @@ def handle_basic_executor_options(options, parser):
 class BaseExecutor(object):
     """Class for starting and handling processes."""
 
-    def __init__(self):
+    def __init__(
+        self,
+        child_setup_fn=None,
+        parent_setup_fn=None,
+        parent_cleanup_fn=None,
+    ):
+        """
+        See BaseExecutor._start_execution for a description of the parameters.
+        """
+        self.child_setup_fn = (
+            child_setup_fn if child_setup_fn is None else util.dummy_fn
+        )
+        self.parent_setup_fn = (
+            parent_setup_fn if parent_setup_fn is None else util.dummy_fn
+        )
+        self.parent_cleanup_fn = (
+            parent_cleanup_fn if parent_cleanup_fn is None else util.dummy_fn
+        )
         self.PROCESS_KILLED = False
         # killing process is triggered asynchronously, need a lock for synchronization
         self.SUB_PROCESS_PIDS_LOCK = threading.Lock()
@@ -75,8 +92,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. It receives one kwarg, tool_pid.
         @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,13 +105,26 @@ def _start_execution(
             and the result of parent_cleanup_fn (do not use os.wait)
         """
 
+        MARKER_PARENT_COMPLETED = b"B"  # noqa: N806 local constant
+
         def pre_subprocess():
             # Do some other setup the caller wants.
             child_setup_fn()
 
             # put us into the cgroup(s)
             pid = os.getpid()
-            cgroups.add_task(pid)
+            if cgroups is not None:
+                cgroups.add_task(pid)
+
+            # Unfortunately, waiting in preexec_fn causes a deadlock
+            # This would otherwise be a great solution
+            #
+            # Wait until parent is also ready
+            # logging.debug("Waiting for parent fn %d", from_parent)
+            # received = os.read(from_parent, 1)
+            # assert received == MARKER_PARENT_COMPLETED, received
+
+            # do exec(args) after this function returns
 
         # Set HOME and TMPDIR to fresh directories.
         tmp_dir = os.path.join(temp_dir, "tmp")
@@ -108,18 +138,52 @@ 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,
-            stdout=stdout,
-            stderr=stderr,
-            env=env,
-            cwd=cwd,
-            close_fds=True,
-            preexec_fn=pre_subprocess,
-        )
+        try:
+            # pipe parent->child
+            from_parent, to_child = os.pipe()
+            logging.debug(
+                "Pipe from parent (read) %d to child (write) %s.", from_parent, to_child
+            )
+
+            script = """
+import sys, os
+os.read(from_parent, 1)
+os.close(from_parent)
+os.execvp(sys.argv[1], sys.argv[1:])
+            """.replace(
+                "from_parent", str(from_parent)
+            )
+
+            # Create process p in a "paused" state
+            p = subprocess.Popen(
+                [sys.executable, "-c", script, *args],
+                stdin=stdin,
+                stdout=stdout,
+                stderr=stderr,
+                env=env,
+                cwd=cwd,
+                close_fds=True,
+                pass_fds=(from_parent,),
+                preexec_fn=pre_subprocess,
+            )
+
+            # Close unnecessary ends of pipes such that read() does not block forever
+            # if all other processes have terminated.
+            os.close(from_parent)
+
+            # Now that we have an active but paused child.
+            # Run the parent_setup_fn
+            logging.debug("Running parent_setup_fn(tool_pid=%d)", p.pid)
+            parent_setup = parent_setup_fn(tool_pid=p.pid)
+
+            # Signal to child that parent is done with setup
+            # It's go time.
+            logging.debug("Signalling child on %d", to_child)
+            os.write(to_child, MARKER_PARENT_COMPLETED)
+        finally:
+            os.close(to_child)
+
+        logging.debug("Waiting on child")
 
         def wait_and_get_result():
             exitcode, ru_child = self._wait_for_process(p.pid, args[0])

benchexec/containerexecutor.py

@@ -346,12 +346,16 @@ def __init__(
         self._uid = (
             uid
             if uid is not None
-            else container.CONTAINER_UID if container_system_config else os.getuid()
+            else container.CONTAINER_UID
+            if container_system_config
+            else os.getuid()
         )
         self._gid = (
             gid
             if gid is not None
-            else container.CONTAINER_GID if container_system_config else os.getgid()
+            else container.CONTAINER_GID
+            if container_system_config
+            else os.getgid()
         )
         self._allow_network = network_access
         self._env_override = {}
@@ -475,9 +479,9 @@ def execute_run(
                 cgroups=cgroups,
                 output_dir=output_dir,
                 result_files_patterns=result_files_patterns,
-                child_setup_fn=util.dummy_fn,
-                parent_setup_fn=util.dummy_fn,
-                parent_cleanup_fn=util.dummy_fn,
+                child_setup_fn=self.child_setup_fn,
+                parent_setup_fn=self.parent_setup_fn,
+                parent_cleanup_fn=self.parent_cleanup_fn,
             )
 
             with self.SUB_PROCESS_PIDS_LOCK:
@@ -647,6 +651,9 @@ def grandchild():
                 # Signal readiness to parent by sending our PID
                 # and wait until parent is also ready
                 os.write(to_parent, str(my_outer_pid).encode())
+                logging.debug(
+                    "Grandchild: Sent my pid; waiting for grandparent to be ready"
+                )
                 received = os.read(from_parent, 1)
                 assert received == MARKER_PARENT_COMPLETED, received
 
@@ -974,11 +981,17 @@ def check_child_exit_code():
             if use_cgroup_ns:
                 cgroups = cgroups.create_fresh_child_cgroup_for_delegation()
 
+            # Do parent setup
+            # Note, we use the grandchild pid rather than child pid
+            # because grandchild is the one actually exec'ing the tool
+            logging.debug("Calling parent_setup_fn(tool_pid=%d)", grandchild_pid)
+            parent_setup = parent_setup_fn(tool_pid=grandchild_pid)
+
             # start measurements
             cgroups.add_task(grandchild_pid)
-            parent_setup = parent_setup_fn()
 
             # Signal grandchild that setup is finished
+            logging.debug("Telling grandchild we are ready")
             os.write(to_grandchild, MARKER_PARENT_COMPLETED)
 
             # Copy file descriptor, otherwise we could not close from_grandchild in

benchexec/runexecutor.py

@@ -829,7 +829,7 @@ 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
             if self._energy_measurement is not None and packages:

benchexec/test_runexecutor.py

@@ -809,6 +809,23 @@ def test_starttime(self):
         self.assertLessEqual(before, run_starttime)
         self.assertLessEqual(run_starttime, after)
 
+    def test_parent_fns(self):
+        if not os.path.exists("/bin/sh"):
+            self.skipTest("missing /bin/sh")
+
+        def parent_setup_fn(*, tool_pid, **kwargs):
+            # I don't want to require psutil just for this
+            # I'll just read the procfs
+            assert os.path.exists("/proc/{tool_pid}")
+            return 12345
+
+        def parent_cleanup_fn(parent_setup, exit_code, path):
+            assert parent_setup == 12345
+            assert exit_code == 123
+
+        self.setUp(parent_setup_fn=parent_setup_fn, parent_cleanup_fn=parent_cleanup_fn)
+        self.execute_run("/bin/sh", "-c", "exit 123")
+
     def test_frozen_process(self):
         # https://github.com/sosy-lab/benchexec/issues/840
         if not os.path.exists(self.sleep):
@@ -1092,6 +1109,23 @@ 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")
+
+        def parent_setup_fn(*, tool_pid, **kwargs):
+            # I don't want to require psutil just for this
+            # I'll just read the procfs
+            assert os.path.exists("/proc/{tool_pid}")
+            return 12345
+
+        def parent_cleanup_fn(parent_setup, exit_code, path):
+            assert parent_setup == 12345
+            assert exit_code == 123
+
+        self.setUp(parent_setup_fn=parent_setup_fn, parent_cleanup_fn=parent_cleanup_fn)
+        self.execute_run("/bin/sh", "-c", "exit 123")
+
     def test_file_count_limit(self):
         if not os.path.exists("/bin/sh"):
             self.skipTest("missing /bin/sh")