Skip to content

Commit

Permalink
Add child_pid kwarg to parent_setup_fn
Browse files Browse the repository at this point in the history
  • Loading branch information
charmoniumQ committed Feb 15, 2024
1 parent 16033e5 commit 9ea65ba
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 10 deletions.
18 changes: 16 additions & 2 deletions benchexec/baseexecutor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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,
Expand All @@ -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])
Expand Down
8 changes: 2 additions & 6 deletions benchexec/containerexecutor.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,16 +346,12 @@ 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 = {}
Expand Down
31 changes: 29 additions & 2 deletions benchexec/test_runexecutor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 9ea65ba

Please sign in to comment.