diff --git a/simpleflow/execute.py b/simpleflow/execute.py index da77e6759..8e7cfd7d2 100644 --- a/simpleflow/execute.py +++ b/simpleflow/execute.py @@ -5,6 +5,8 @@ import time +import psutil + try: import subprocess32 as subprocess except ImportError: @@ -172,7 +174,7 @@ def wait_subprocess(process, timeout=None, command_info=None): return process.wait() -def python(interpreter='python', logger_name=__name__, timeout=None): +def python(interpreter='python', logger_name=__name__, timeout=None, kill_children=False): """ Execute a callable as an external Python program. @@ -202,6 +204,8 @@ def execute(*args, **kwargs): '--result-fd={}'.format(dup_result_fd), '--error-fd={}'.format(dup_error_fd), ] + if kill_children: + full_command.append('--kill-children') if compat.PY2: # close_fds doesn't work with python2 (using its C _posixsubprocess helper) close_fds = False pass_fds = () @@ -214,9 +218,7 @@ def execute(*args, **kwargs): close_fds=close_fds, pass_fds=pass_fds, ) - rc = wait_subprocess(process, timeout=timeout, command_info=full_command) - os.close(dup_result_fd) os.close(dup_error_fd) if rc: @@ -409,7 +411,29 @@ def main(): metavar='N', help='error file descriptor', ) + parser.add_argument( + '--kill-children', + action='store_true', + help='kill child processes on exit', + ) cmd_arguments = parser.parse_args() + + def kill_child_processes(): + process = psutil.Process(os.getpid()) + children = process.children(recursive=True) + + for child in children: + try: + child.terminate() + except psutil.NoSuchProcess: + pass + _, still_alive = psutil.wait_procs(children, timeout=0.3) + for child in still_alive: + try: + child.kill() + except psutil.NoSuchProcess: + pass + funcname = cmd_arguments.funcname try: arguments = format.decode(cmd_arguments.funcargs) @@ -447,6 +471,8 @@ def main(): if not compat.PY2: details = details.encode('utf-8') os.write(cmd_arguments.error_fd, details) + if cmd_arguments.kill_children: + kill_child_processes() sys.exit(1) if cmd_arguments.result_fd == 1: # stdout (legacy) @@ -456,6 +482,8 @@ def main(): if not compat.PY2: result = result.encode('utf-8') os.write(cmd_arguments.result_fd, result) + if cmd_arguments.kill_children: + kill_child_processes() if __name__ == '__main__': diff --git a/tests/test_simpleflow/test_execute.py b/tests/test_simpleflow/test_execute.py index 985efc5ab..06fba083c 100644 --- a/tests/test_simpleflow/test_execute.py +++ b/tests/test_simpleflow/test_execute.py @@ -7,9 +7,12 @@ import platform import threading +import psutil import pytest import time +import subprocess + from simpleflow import execute from simpleflow.exceptions import ExecutionError, ExecutionTimeoutError @@ -287,3 +290,21 @@ def test_timeout_execute_from_thread(): t = threading.Thread(target=test_timeout_execute) t.start() t.join() + + +def create_sleeper_subprocess(): + pid = subprocess.Popen(['sleep', '600']).pid + return pid + + +def test_execute_dont_kill_children(): + pid = execute.python()(create_sleeper_subprocess)() + subprocess = psutil.Process(pid) + assert subprocess.status() == 'sleeping' + subprocess.terminate() # cleanup + + +def test_execute_kill_children(): + pid = execute.python(kill_children=True)(create_sleeper_subprocess)() + with pytest.raises(psutil.NoSuchProcess): + psutil.Process(pid)