diff --git a/pmb/helpers/run_core.py b/pmb/helpers/run_core.py index 9b97ac6e3c602addbea60ffdc4fecf9d32b93282..f3b4cb38aa61dbe0d230fb8f66d1ade80daab222 100644 --- a/pmb/helpers/run_core.py +++ b/pmb/helpers/run_core.py @@ -93,6 +93,45 @@ def pipe_read(args, process, output_to_stdout=False, output_return=False, return +def kill_process_tree(args, pid, ppids, kill_as_root): + """ + Recursively kill a pid and its child processes + + :param pid: process id that will be killed + :param ppids: list of process id and parent process id tuples (pid, ppids) + :param kill_as_root: use sudo to kill the process + """ + if kill_as_root: + pmb.helpers.run.root(args, ["kill", "-9", str(pid)], + check=False) + else: + pmb.helpers.run.user(args, ["kill", "-9", str(pid)], + check=False) + + for (child_pid, child_ppid) in ppids: + if child_ppid == str(pid): + kill_process_tree(args, child_pid, ppids, kill_as_root) + + +def kill_command(args, pid, kill_as_root): + """ + Kill a command process and recursively kill its child processes + + :param pid: process id that will be killed + :param kill_as_root: use sudo to kill the process + """ + cmd = ["ps", "-e", "-o", "pid=,ppid=", "--noheaders"] + ret = subprocess.run(cmd, check=True, stdout=subprocess.PIPE) + ppids = [] + proc_entries = ret.stdout.decode("utf-8").split('\n') + for row in proc_entries: + items = row.split() + if len(items) == 2: + ppids.append(items) + + kill_process_tree(args, pid, ppids, kill_as_root) + + def foreground_pipe(args, cmd, working_dir=None, output_to_stdout=False, output_return=False, output_timeout=True, kill_as_root=False): @@ -141,11 +180,7 @@ def foreground_pipe(args, cmd, working_dir=None, output_to_stdout=False, str(args.timeout) + " seconds. Killing it.") logging.info("NOTE: The timeout can be increased with" " 'pmbootstrap -t'.") - if kill_as_root: - pmb.helpers.run.root(args, ["kill", "-9", - str(process.pid)]) - else: - process.kill() + kill_command(args, process.pid, kill_as_root) continue # Read all currently available output diff --git a/test/test_run_core.py b/test/test_run_core.py index a3b0b16c1c67a8071ee41e6d6e243298fe3e1ae2..6afd56669ba83b4bc3d508fe61cc6ca107f9e1d8 100644 --- a/test/test_run_core.py +++ b/test/test_run_core.py @@ -23,6 +23,7 @@ This file tests functions from pmb.helpers.run_core import os import sys +import subprocess import pytest # Import from parent directory @@ -108,6 +109,27 @@ def test_foreground_pipe(args): ret = func(args, cmd, output_return=True, output_timeout=True) assert ret == (0, "first\nsecond\nthird\nfourth\n") + # Check if all child processes are killed after timeout. + cmd = ["sudo", "sh", "-c", + "pgid=$(ps -p ${1:-$$} -o pgid=);echo $pgid | tr -d '\n';" + + "sleep 10 | sleep 20 | sleep 30"] + args.timeout = 0.3 + ret = func(args, cmd, output_return=True, output_timeout=True, + kill_as_root=True) + pgid = str(ret[1]) + + cmd = ["ps", "-e", "-o", "pgid=,comm=", "--noheaders"] + ret = subprocess.run(cmd, check=True, stdout=subprocess.PIPE) + procs = str(ret.stdout.decode("utf-8")).split('\n') + child_procs = [] + for process in procs: + items = process.split() + if len(items) != 2: + continue + if pgid == items[0] and "sleep" in items[1]: + child_procs.append(items) + assert len(child_procs) == 0 + def test_foreground_tui(): func = pmb.helpers.run_core.foreground_tui