Add child_pid kwarg to parent_setup_fn

benchexec/baseexecutor.py

@@ -8,6 +8,7 @@
 import errno
 import logging
 import os
+import select
 import subprocess
 import sys
 import threading
@@ -88,14 +89,23 @@ 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(to_child)
+            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 +118,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 +126,18 @@ def pre_subprocess():
             env=env,
             cwd=cwd,
             close_fds=True,
+            pass_fds=(to_parent, from_parent),
             preexec_fn=pre_subprocess,
         )
+        print("Continuing")
+
+        os.close(to_parent)
+        os.close(from_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/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):
@@ -1101,15 +1130,13 @@ def test_parent_fns(self):
         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
-            print("parent setup fn")
             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):
-            print("parent cleanup fn")
             assert parent_setup == 12345
             nonlocal parent_cleanup_ran
             parent_cleanup_ran = True