feat(run-task): send SIGTERM to the task's processes shortly before maxRunTime

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
 

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"

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()

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"