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,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__":

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,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

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"