diff --git a/src/taskgraph/run-task/run-task b/src/taskgraph/run-task/run-task index 4e4e28299..c8b7a132e 100755 --- a/src/taskgraph/run-task/run-task +++ b/src/taskgraph/run-task/run-task @@ -22,6 +22,7 @@ import json import os import platform import re +import runpy import shutil import socket import stat @@ -30,6 +31,7 @@ import sys import tempfile import threading import time +import traceback import urllib.error import urllib.request from pathlib import Path @@ -112,16 +114,58 @@ IS_WINDOWS = os.name == "nt" NULL_REVISION = "0000000000000000000000000000000000000000" -def print_line(prefix, m): +def format_line(prefix, m): now = ( datetime.datetime.now(tz=datetime.timezone.utc) .isoformat(timespec="milliseconds") .encode("utf-8") ) - sys.stdout.buffer.write(b"[%s %s] %s" % (prefix, now, m)) + return b"[%s %s] %s" % (prefix, now, m) + + +def print_line(prefix, m): + sys.stdout.buffer.write(format_line(prefix, m)) sys.stdout.buffer.flush() +class _PrefixedLineWriter: + """File-like object that prefixes each complete line written to it + with the current time, before forwarding it to the given buffer.""" + + def __init__(self, prefix, buffer): + self.prefix = prefix + self.buffer = buffer + self._pending = "" + + def write(self, s): + self._pending += s + *lines, self._pending = self._pending.split("\n") + for line in lines: + self.buffer.write(format_line(self.prefix, b"%s\n" % line.encode("utf-8"))) + self.buffer.flush() + return len(s) + + def flush(self): + if self._pending: + self.buffer.write(format_line(self.prefix, self._pending.encode("utf-8"))) + self._pending = "" + self.buffer.flush() + + +@contextlib.contextmanager +def prefix_output(prefix): + """Redirect stdout/stderr so each line written is prefixed with the + current time, like output from `run_command`.""" + writer = _PrefixedLineWriter(prefix, sys.stdout.buffer) + old_stdout, old_stderr = sys.stdout, sys.stderr + sys.stdout = sys.stderr = writer + try: + yield + finally: + writer.flush() + sys.stdout, sys.stderr = old_stdout, old_stderr + + def reap_zombies(main_subprocess): """Wait for main_subprocess to exit, while awaiting any other child processes""" while main_subprocess.poll() is None: @@ -319,6 +363,24 @@ def run_command(prefix, args, *, extra_env=None, cwd=None, stdin_data=None): return p.wait() +def run_python_script(path): + """Execute a Python script in-process.""" + print_line( + b"setup", + b"running python script: %s\n" % path.encode("utf-8"), + ) + try: + with prefix_output(b"script"): + runpy.run_path(path) + except Exception: + print_line( + b"setup", + b"script %s failed:\n%s" + % (path.encode("utf-8"), traceback.format_exc().encode("utf-8")), + ) + sys.exit(1) + + def get_posix_user_group(user, group): import grp # noqa: PLC0415 import pwd # noqa: PLC0415 @@ -1668,6 +1730,7 @@ def main(args): "MOZ_FETCHES_DIR", "MOZ_PYTHON_HOME", "MOZ_UV_HOME", + "RUN_TASK_PRE_COMMAND_HOOK", "PIP_CACHE_DIR", "UPLOAD_DIR", "UV_CACHE_DIR", @@ -1691,6 +1754,9 @@ def main(args): # fetches to grab dependencies. install_pip_requirements(repositories) + 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) diff --git a/test/test_scripts_run_task.py b/test/test_scripts_run_task.py index 57420be14..8b184178f 100644 --- a/test/test_scripts_run_task.py +++ b/test/test_scripts_run_task.py @@ -761,6 +761,55 @@ def test_main_abspath_environment(mocker, run_main): assert env[key] == "/builds/worker/file" +def test_pre_task_run_hook_sets_env(run_main, tmp_path): + hook = tmp_path / "hook.py" + hook.write_text("import os\nos.environ['HOOK_RAN'] = '1'\n") + + result, env = run_main(env={"RUN_TASK_PRE_COMMAND_HOOK": str(hook)}) + + assert result == 0 + assert env.get("HOOK_RAN") == "1" + + +def test_pre_task_run_hook_failure_aborts_before_task( + run_main, patch_run_command, tmp_path, capsys +): + called_with = patch_run_command() + hook = tmp_path / "hook.py" + hook.write_text("raise RuntimeError('boom')") + + with pytest.raises(SystemExit) as excinfo: + run_main(env={"RUN_TASK_PRE_COMMAND_HOOK": str(hook)}) + + assert excinfo.value.code == 1 + assert called_with == [] + + output = capsys.readouterr().out + assert "script" in output and "failed" in output + assert "RuntimeError: boom" in output + assert "hook.py" in output and "line 1" in output + + +def test_pre_task_run_hook_output_is_prefixed(run_main, tmp_path, capsys): + hook = tmp_path / "hook.py" + hook.write_text("print('hello')\nprint('world')\n") + + result, env = run_main(env={"RUN_TASK_PRE_COMMAND_HOOK": str(hook)}) + assert result == 0 + + lines = [ + line + for line in capsys.readouterr().out.splitlines() + if line.startswith("[script ") and line.endswith(("] hello", "] world")) + ] + assert len(lines) == 2 + + +def test_no_pre_task_run_hook_is_noop(run_main): + result, env = run_main(env={}) + assert result == 0 + + SPARSE_REPO_FILES = [ "a/deep/two.txt", "a/one.txt",