diff --git a/src/taskgraph/run-task/run-task b/src/taskgraph/run-task/run-task index 8508318af..22c40c2bc 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,88 @@ 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. +max_run_time_abort = threading.Event() + + +def ignore_sigterm(signum, frame): + """Ignores the SIGTERM that abort_task_processes() sends to the process + group of run-task, so that run-task keeps forwarding the output of its + command until it exits. + + A handler rather than SIG_IGN, which the task's processes would inherit + across exec and then never get the 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 holds exactly + the task's processes: generic-worker starts run-task in a new process + group, and in docker it is PID 1. The worker later kills that group as a + whole. run-task ignores the signal itself, see ignore_sigterm(). + """ + 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 %ds early\n" % (max_run_time, margin), + ) + try: + os.killpg(os.getpgrp(), 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. It must + be called from the main thread, where the SIGTERM handler is installed. + + Nothing is scheduled on Windows, where the worker kills the task's + processes. + """ + 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 + signal.signal(signal.SIGTERM, ignore_sigterm) + timer = threading.Timer( + max_run_time - margin, abort_task_processes, (max_run_time, margin) + ) + timer.daemon = True + timer.start() + return timer + + 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 @@ -291,7 +374,7 @@ def retry_required_command(prefix, args, *, extra_env=None, cwd=None, retries=2) res = run_command(prefix, args, extra_env=extra_env, cwd=cwd) if not res: return - if not retries: + if not retries or max_run_time_abort.is_set(): sys.exit(res) retries -= 1 backoff *= 2 @@ -301,8 +384,18 @@ def retry_required_command(prefix, args, *, extra_env=None, cwd=None, retries=2) 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. + Returns the process exit code, or 1 once the task has been aborted for + exceeding its maxRunTime, see abort_task_processes(). A process that + hasn't started by then isn't started at all. """ + if max_run_time_abort.is_set(): + print_line( + prefix, + b"run-task: the task was aborted for exceeding its maxRunTime, " + b"not executing %r\n" % (args,), + ) + return 1 + print_line(prefix, b"executing %r\n" % args) env = dict(os.environ) @@ -360,7 +453,15 @@ 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() + if max_run_time_abort.is_set(): + print_line( + prefix, + b"run-task: the task was aborted for exceeding its maxRunTime, " + b"exiting with code 1\n", + ) + return 1 + return returncode def run_python_script(path): @@ -1513,6 +1614,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 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..06ca8418b 100644 --- a/test/test_scripts_run_task.py +++ b/test/test_scripts_run_task.py @@ -1,11 +1,14 @@ import functools import json import os +import signal import site import stat import subprocess import sys import tempfile +import threading +import time from argparse import Namespace from importlib.machinery import SourceFileLoader from importlib.util import module_from_spec, spec_from_loader @@ -1153,3 +1156,155 @@ def test_git_checkout_sparse_cone_converts_for_globs( after = materialized_files(destination) assert set(before) <= set(after) assert "c/four.md" in after + + +@pytest.fixture +def sigterm_handler(): + previous = signal.getsignal(signal.SIGTERM) + yield previous + signal.signal(signal.SIGTERM, previous) + + +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, sigterm_handler, 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 + assert signal.getsignal(signal.SIGTERM) is sigterm_handler + return + try: + assert timer.interval == max_run_time - margin + assert timer.daemon + assert timer.args == (max_run_time, margin) + assert signal.getsignal(signal.SIGTERM) is run_task_mod.ignore_sigterm + finally: + timer.cancel() + + +@nowin +def test_ignore_sigterm_not_inherited(run_task_mod, sigterm_handler): + signal.signal(signal.SIGTERM, run_task_mod.ignore_sigterm) + os.kill(os.getpid(), signal.SIGTERM) + # Unlike SIG_IGN, the handler is not inherited across exec. + child = subprocess.run(["sh", "-c", "kill -TERM $$"]) + assert child.returncode == -signal.SIGTERM + + +@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, "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_run_command_keeps_forwarding_output_after_abort( + mocker, mock_stdin, run_task_mod, abort_flag, capsys, tmp_path +): + # Signal the command alone rather than the process group of the test + # runner. + pid_file = tmp_path / "pid" + + def kill_command(pgid, sig): + for _ in range(100): + if pid_file.exists() and pid_file.read_text().strip(): + break + time.sleep(0.05) + os.kill(int(pid_file.read_text()), sig) + + mocker.patch.object(run_task_mod.os, "killpg", side_effect=kill_command) + 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", + f"echo $$ > '{pid_file}'; " + "trap 'echo got TERM; kill $!; exit 3' TERM; " + "sleep 30 >/dev/null 2>&1 & wait $!; echo not aborted", + ], + ) + finally: + timer.cancel() + + assert returncode == 1 + 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 any("exiting with code 1" in line for line in output) + assert not any("not aborted" in line for line in output) + + +def test_run_command_not_started_after_abort( + run_task_mod, abort_flag, capsys, tmp_path +): + abort_flag.set() + marker = tmp_path / "started" + + returncode = run_task_mod.run_command( + b"task", [sys.executable, "-c", f"open({str(marker)!r}, 'w')"] + ) + + assert returncode == 1 + assert not marker.exists() + assert "not executing" in capsys.readouterr().out + + +def test_retry_required_command_not_retried_after_abort( + mocker, run_task_mod, abort_flag +): + abort_flag.set() + sleep = mocker.patch.object(run_task_mod.time, "sleep") + + with pytest.raises(SystemExit) as excinfo: + run_task_mod.retry_required_command(b"vcs", [sys.executable, "-c", "pass"]) + + assert excinfo.value.code == 1 + sleep.assert_not_called() 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"