diff --git a/src/taskgraph/run-task/run-task b/src/taskgraph/run-task/run-task index 8508318af..ac88ed1c2 100755 --- a/src/taskgraph/run-task/run-task +++ b/src/taskgraph/run-task/run-task @@ -24,6 +24,7 @@ import platform import re import runpy import shutil +import signal import socket import stat import subprocess @@ -177,6 +178,118 @@ def reap_zombies(main_subprocess): main_subprocess._handle_exitstatus(status) +# Shortly before a task exceeds its maxRunTime, run-task sends SIGTERM to the +# task's processes, so that they can shut down and flush their logs before the +# worker kills them. The SIGTERM is sent 5% of maxRunTime early, and at most +# this many seconds early. +MAX_RUN_TIME_ABORT_MARGIN = 60 + +# Set once run-task has sent SIGTERM to the task's processes itself. +max_run_time_abort = threading.Event() + +# The process currently run by run_command(), if any. +current_child = None + + +@contextlib.contextmanager +def deferred_sigterm(): + """Defers SIGTERM until the end of the context. + + Shortly before a task exceeds its maxRunTime, run-task sends SIGTERM to all + of the task's processes, itself included, so that they can shut down before + the worker kills them. This lets the process run in the context shut down + while its output is still being forwarded. + + If the SIGTERM was sent by run-task itself (see abort_task_processes()), + it is not re-raised afterwards, so that run-task exits normally with the + exit code of its command. + """ + if not IS_POSIX or threading.current_thread() is not threading.main_thread(): + yield + return + received = [] + previous = signal.signal(signal.SIGTERM, lambda signum, frame: received.append(1)) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + if received and not max_run_time_abort.is_set(): + os.kill(os.getpid(), signal.SIGTERM) + + +def abort_task_processes(max_run_time, margin): + """Sends SIGTERM to the task's processes, so that they can shut down and + flush their logs before the worker kills them when the task exceeds its + maxRunTime. + + The signal is sent to the process group of run-task, which is its leader + under generic-worker, and PID 1 in docker, where the signal then reaches + every process in the container except run-task itself. If run-task is + neither, e.g. when run through a shell, only its current command is + signalled, so that the shell keeps waiting for run-task. + """ + max_run_time_abort.set() + print_line( + b"task", + b"run-task: the task is about to exceed its maxRunTime of %ds; sending " + b"SIGTERM to the task's processes %ds early so that they can shut down " + b"before the worker kills them\n" % (max_run_time, margin), + ) + try: + if os.getpid() in (1, os.getpgrp()): + os.killpg(os.getpgrp(), signal.SIGTERM) + elif current_child is not None: + current_child.send_signal(signal.SIGTERM) + except OSError as e: + print_line( + b"task", b"run-task: could not send SIGTERM: %s\n" % str(e).encode("utf-8") + ) + + +def schedule_max_run_time_abort(): + """Arranges for abort_task_processes() to be called shortly before the + task exceeds its maxRunTime, which taskgraph passes in the + TASKCLUSTER_MAX_RUN_TIME environment variable. + + The abort happens 5% of maxRunTime early, at most MAX_RUN_TIME_ABORT_MARGIN + seconds early. That is measured from now, so this must be called as soon + as run-task starts, which is also when the worker starts counting. + + On Windows, the task's processes are killed by the worker, as before. + """ + value = os.environ.get("TASKCLUSTER_MAX_RUN_TIME") + if not value or not IS_POSIX: + return None + try: + max_run_time = int(value) + except ValueError: + max_run_time = 0 + if max_run_time <= 0: + print_line( + b"setup", + b"ignoring invalid TASKCLUSTER_MAX_RUN_TIME: %s\n" % value.encode("utf-8"), + ) + return None + margin = min(MAX_RUN_TIME_ABORT_MARGIN, max_run_time // 20) + if margin < 1: + return None + timer = threading.Timer( + max_run_time - margin, abort_task_processes, (max_run_time, margin) + ) + timer.daemon = True + timer.start() + return timer + + +def exit_code_after_abort(returncode): + """Returns the exit code of run-task after it aborted the task's processes + itself: the exit code of its command, reported as a shell would when the + command was killed by a signal, and never 0, since the task was aborted.""" + if returncode < 0: + returncode = 128 - returncode + return returncode or 1 + + def _call_windows_retry(func, args=(), retry_max=5, retry_delay=0.5): """ It's possible to see spurious errors on Windows due to various things @@ -298,11 +411,14 @@ def retry_required_command(prefix, args, *, extra_env=None, cwd=None, retries=2) time.sleep(backoff) +@deferred_sigterm() def run_command(prefix, args, *, extra_env=None, cwd=None, stdin_data=None): """Runs a process and prefixes its output with the time. Returns the process exit code. """ + global current_child + print_line(prefix, b"executing %r\n" % args) env = dict(os.environ) @@ -341,6 +457,8 @@ def run_command(prefix, args, *, extra_env=None, cwd=None, stdin_data=None): if stdin_file is not None: stdin_file.close() + current_child = p + stdout = io.TextIOWrapper(p.stdout, encoding="latin1") if os.getpid() == 1: @@ -360,7 +478,9 @@ def run_command(prefix, args, *, extra_env=None, cwd=None, stdin_data=None): if reaper_thread: reaper_thread.join() - return p.wait() + returncode = p.wait() + current_child = None + return returncode def run_python_script(path): @@ -1513,6 +1633,7 @@ def main(args): b"setup", b"Invoked by command: %s\n" % " ".join(args).encode("utf-8"), ) + schedule_max_run_time_abort() _display_python_versions() running_as_root = IS_POSIX and os.getuid() == 0 @@ -1760,7 +1881,16 @@ def main(args): if hook_path := os.environ.get("RUN_TASK_PRE_COMMAND_HOOK"): run_python_script(hook_path) - return run_command(b"task", task_args, cwd=args.task_cwd) + returncode = run_command(b"task", task_args, cwd=args.task_cwd) + + if max_run_time_abort.is_set(): + returncode = exit_code_after_abort(returncode) + print_line( + b"task", + b"run-task: the task was aborted for exceeding its maxRunTime, " + b"exiting with code %d\n" % returncode, + ) + return returncode if __name__ == "__main__": diff --git a/src/taskgraph/transforms/task.py b/src/taskgraph/transforms/task.py index 1cbd96136..97dbfd809 100644 --- a/src/taskgraph/transforms/task.py +++ b/src/taskgraph/transforms/task.py @@ -379,6 +379,9 @@ def build_docker_worker_payload(config, task, task_def): if "max-run-time" in worker: payload["maxRunTime"] = worker["max-run-time"] + # Let run-task know how long the task may run, so that it can abort + # the task's processes shortly before the worker kills them. + payload["env"]["TASKCLUSTER_MAX_RUN_TIME"] = str(worker["max-run-time"]) run_task = os.path.basename(payload.get("command", [""])[0]).startswith("run-task") @@ -583,6 +586,9 @@ def build_generic_worker_payload(config, task, task_def): task_def["payload"]["onExitStatus"] = on_exit_status env = worker.get("env", {}) + # Let run-task know how long the task may run, so that it can abort the + # task's processes shortly before the worker kills them. + env["TASKCLUSTER_MAX_RUN_TIME"] = str(worker["max-run-time"]) if task.get("needs-sccache"): env["USE_SCCACHE"] = "1" diff --git a/test/test_scripts_run_task.py b/test/test_scripts_run_task.py index 86a6a8c44..3ce26c9d2 100644 --- a/test/test_scripts_run_task.py +++ b/test/test_scripts_run_task.py @@ -1,11 +1,13 @@ import functools import json import os +import signal import site import stat import subprocess import sys import tempfile +import threading from argparse import Namespace from importlib.machinery import SourceFileLoader from importlib.util import module_from_spec, spec_from_loader @@ -1153,3 +1155,124 @@ def test_git_checkout_sparse_cone_converts_for_globs( after = materialized_files(destination) assert set(before) <= set(after) assert "c/four.md" in after + + +def test_schedule_max_run_time_abort_without_env(monkeypatch, run_task_mod): + monkeypatch.delenv("TASKCLUSTER_MAX_RUN_TIME", raising=False) + assert run_task_mod.schedule_max_run_time_abort() is None + + +@nowin +@pytest.mark.parametrize("value", ("", "soon", "0", "-5")) +def test_schedule_max_run_time_abort_invalid_value( + monkeypatch, run_task_mod, capsys, value +): + monkeypatch.setenv("TASKCLUSTER_MAX_RUN_TIME", value) + assert run_task_mod.schedule_max_run_time_abort() is None + out = capsys.readouterr().out + assert ("ignoring invalid TASKCLUSTER_MAX_RUN_TIME" in out) == bool(value) + + +@nowin +@pytest.mark.parametrize( + "max_run_time,margin", + ( + (3600, 60), + (1200, 60), + (600, 30), + (100, 5), + (20, 1), + (19, None), + ), +) +def test_schedule_max_run_time_abort_margin( + monkeypatch, run_task_mod, max_run_time, margin +): + monkeypatch.setenv("TASKCLUSTER_MAX_RUN_TIME", str(max_run_time)) + timer = run_task_mod.schedule_max_run_time_abort() + if margin is None: + assert timer is None + return + try: + assert timer.interval == max_run_time - margin + assert timer.daemon + assert timer.args == (max_run_time, margin) + finally: + timer.cancel() + + +@pytest.fixture +def abort_flag(run_task_mod): + run_task_mod.max_run_time_abort.clear() + yield run_task_mod.max_run_time_abort + run_task_mod.max_run_time_abort.clear() + + +@nowin +def test_abort_task_processes_signals_process_group( + mocker, run_task_mod, abort_flag, capsys +): + mocker.patch.object(run_task_mod.os, "getpid", return_value=1234) + mocker.patch.object(run_task_mod.os, "getpgrp", return_value=1234) + killpg = mocker.patch.object(run_task_mod.os, "killpg") + + run_task_mod.abort_task_processes(600, 30) + + killpg.assert_called_once_with(1234, signal.SIGTERM) + assert abort_flag.is_set() + assert "maxRunTime of 600s" in capsys.readouterr().out + + +@nowin +def test_abort_task_processes_signals_child_when_not_group_leader( + mocker, run_task_mod, abort_flag +): + mocker.patch.object(run_task_mod.os, "getpid", return_value=1234) + mocker.patch.object(run_task_mod.os, "getpgrp", return_value=1000) + killpg = mocker.patch.object(run_task_mod.os, "killpg") + child = Mock() + mocker.patch.object(run_task_mod, "current_child", child) + + run_task_mod.abort_task_processes(600, 30) + + killpg.assert_not_called() + child.send_signal.assert_called_once_with(signal.SIGTERM) + assert abort_flag.is_set() + + +@nowin +def test_run_command_keeps_forwarding_output_after_abort( + mocker, mock_stdin, run_task_mod, abort_flag, capsys +): + # Make abort_task_processes() signal only the child, so that the process + # group of the test runner is left alone. + mocker.patch.object(run_task_mod.os, "getpgrp", return_value=-1) + timer = threading.Timer(0.5, run_task_mod.abort_task_processes, (20, 1)) + timer.start() + try: + returncode = run_task_mod.run_command( + b"task", + [ + "sh", + "-c", + "trap 'echo got TERM; kill $!; exit 3' TERM; " + "sleep 30 >/dev/null 2>&1 & wait $!; echo not aborted", + ], + ) + finally: + timer.cancel() + + assert returncode == 3 + assert abort_flag.is_set() + # Skip the line echoing the command itself. + output = capsys.readouterr().out.splitlines()[1:] + assert any("got TERM" in line for line in output) + assert not any("not aborted" in line for line in output) + + +@pytest.mark.parametrize( + "returncode,expected", + ((0, 1), (3, 3), (-15, 143), (-9, 137)), +) +def test_exit_code_after_abort(run_task_mod, returncode, expected): + assert run_task_mod.exit_code_after_abort(returncode) == expected diff --git a/test/test_transforms_task.py b/test/test_transforms_task.py index 3d49e7998..5bdb107ab 100644 --- a/test/test_transforms_task.py +++ b/test/test_transforms_task.py @@ -1000,3 +1000,44 @@ def test_dict_schema_accepts_valid_payload(dict_schema_builder, payload): def test_dict_schema_rejects_invalid_payload(dict_schema_builder, payload): with pytest.raises(voluptuous.MultipleInvalid): dict_schema_builder(payload) + + +@pytest.mark.parametrize( + "worker_type,worker", + ( + pytest.param( + "t-linux", + { + "docker-image": "fake-image-name", + "max-run-time": 1800, + "command": ["run-task"], + }, + id="docker-worker", + ), + pytest.param( + "t-win", + { + "max-run-time": 1800, + "command": ["run-task"], + }, + id="generic-worker", + ), + ), +) +def test_max_run_time_passed_to_task_env( + run_transform, make_transform_config, worker_type, worker +): + """run-task is told the maxRunTime of the task through its environment.""" + task_def = { + "description": "fake description", + "name": "fake-task-name", + "worker-type": worker_type, + "worker": worker, + } + + task_dict = run_transform( + task.transforms, task_def, config=make_transform_config() + )[0] + payload = task_dict["task"]["payload"] + assert payload["maxRunTime"] == 1800 + assert payload["env"]["TASKCLUSTER_MAX_RUN_TIME"] == "1800"