feat(run-task): implement hook point before task payload

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)
 
 

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