diff --git a/src/taskgraph/run-task/fetch-content b/src/taskgraph/run-task/fetch-content index 6adb2efbf..e63e2e5c0 100755 --- a/src/taskgraph/run-task/fetch-content +++ b/src/taskgraph/run-task/fetch-content @@ -57,6 +57,50 @@ def log(msg): sys.stderr.flush() +def open_for_append(path): + """Open path for writes that are each appended atomically, even when + several processes append to it concurrently.""" + if sys.platform != "win32": + return os.open(path, os.O_WRONLY | os.O_APPEND) + + import _winapi + import msvcrt + + # O_APPEND seeks to the end before each write on Windows, which isn't + # atomic. Writes through a handle that may only append go to the end of + # the file atomically. + FILE_APPEND_DATA = 0x4 + FILE_SHARE_READ_WRITE_DELETE = 0x7 + handle = _winapi.CreateFile( + path, + FILE_APPEND_DATA, + FILE_SHARE_READ_WRITE_DELETE, + 0, + _winapi.OPEN_EXISTING, + 0, + 0, + ) + return msvcrt.open_osfhandle(handle, 0) + + +def record_marker(name, start, end, text): + """Add a marker to the profile of the resource monitor run-task runs, if + any, using the forwarding protocol of mozsystemmonitor.""" + path = os.environ.get("MOZ_RESOURCE_MONITOR_FORWARD") + if not path: + return + # 2 is the "Tasks" category. + record = ["marker", name, start, end, {"type": "Text", "text": text}, 2] + try: + fd = open_for_append(path) + try: + os.write(fd, (json.dumps(record) + "\n").encode("utf-8")) + finally: + os.close(fd) + except OSError: + pass + + class IntegrityError(Exception): """Represents an integrity error when downloading a URL.""" @@ -971,6 +1015,21 @@ def fetch_and_extract(url, dest_dir, extract=True, sha256=None, size=None): pass +@contextlib.contextmanager +def fetch_marker(url): + """Record the duration of a fetch, including failures.""" + start = time.monotonic() + try: + yield + finally: + record_marker("fetch-content", start, time.monotonic(), url) + + +def fetch_and_extract_with_marker(url, *args): + with fetch_marker(url): + fetch_and_extract(url, *args) + + def fetch_urls(downloads): """Fetch URLs pairs to a pathlib.Path. @@ -984,7 +1043,7 @@ def fetch_urls(downloads): extractions = [] for url, dest_dir, *rest in downloads: staging_dir = pathlib.Path(tempfile.mkdtemp(prefix=".fetch.", dir=dest_dir)) - future = e.submit(fetch_and_extract, url, staging_dir, *rest) + future = e.submit(fetch_and_extract_with_marker, url, staging_dir, *rest) extractions.append((future, staging_dir, dest_dir)) for future, staging_dir, dest_dir in extractions: @@ -1154,80 +1213,82 @@ def git_checkout_archive( def command_git_checkout_archive(args): - dest = pathlib.Path(args.dest) + with fetch_marker(args.repo): + dest = pathlib.Path(args.dest) - try: - git_checkout_archive( - dest, - args.repo, - args.commit, - prefix=args.path_prefix, - ssh_key=args.ssh_key_secret, - include_dot_git=args.include_dot_git, - ) - except Exception: try: - dest.unlink() - except FileNotFoundError: - pass + git_checkout_archive( + dest, + args.repo, + args.commit, + prefix=args.path_prefix, + ssh_key=args.ssh_key_secret, + include_dot_git=args.include_dot_git, + ) + except Exception: + try: + dest.unlink() + except FileNotFoundError: + pass - raise + raise def command_static_url(args): - gpg_sig_url = args.gpg_sig_url - gpg_env_key = args.gpg_key_env - - if bool(gpg_sig_url) != bool(gpg_env_key): - print("--gpg-sig-url and --gpg-key-env must both be defined") - return 1 + with fetch_marker(args.url): + gpg_sig_url = args.gpg_sig_url + gpg_env_key = args.gpg_key_env - if gpg_sig_url: - gpg_signature = b"".join(stream_download(gpg_sig_url)) - gpg_key = os.environb[gpg_env_key.encode("ascii")] + if bool(gpg_sig_url) != bool(gpg_env_key): + print("--gpg-sig-url and --gpg-key-env must both be defined") + return 1 - dest = pathlib.Path(args.dest) - dest.parent.mkdir(parents=True, exist_ok=True) + if gpg_sig_url: + gpg_signature = b"".join(stream_download(gpg_sig_url)) + gpg_key = os.environb[gpg_env_key.encode("ascii")] - basename = urllib.parse.urlparse(args.url).path.split("/")[-1] - if basename.endswith("".join(dest.suffixes)): - dl_dest = dest - else: - dl_dest = dest.parent / basename + dest = pathlib.Path(args.dest) + dest.parent.mkdir(parents=True, exist_ok=True) - try: - download_to_path( - args.url, dl_dest, sha256=args.sha256, size=args.size, headers=args.headers - ) - - if gpg_sig_url: - gpg_verify_path(dl_dest, gpg_key, gpg_signature) + basename = urllib.parse.urlparse(args.url).path.split("/")[-1] + if basename.endswith("".join(dest.suffixes)): + dl_dest = dest + else: + dl_dest = dest.parent / basename - if args.force_archive or should_repack_archive( - dl_dest, dest, args.strip_components, args.add_prefix - ): - repack_archive( - dl_dest, - dest, - args.strip_components, - args.add_prefix, - args.force_archive, + try: + download_to_path( + args.url, dl_dest, sha256=args.sha256, size=args.size, headers=args.headers ) - elif dl_dest != dest: - log(f"Renaming {dl_dest} to {dest}") - dl_dest.rename(dest) - except Exception: - try: - dl_dest.unlink() - except FileNotFoundError: - pass + if gpg_sig_url: + gpg_verify_path(dl_dest, gpg_key, gpg_signature) + + if args.force_archive or should_repack_archive( + dl_dest, dest, args.strip_components, args.add_prefix + ): + repack_archive( + dl_dest, + dest, + args.strip_components, + args.add_prefix, + args.force_archive, + ) + elif dl_dest != dest: + log(f"Renaming {dl_dest} to {dest}") + dl_dest.rename(dest) - raise + except Exception: + try: + dl_dest.unlink() + except FileNotFoundError: + pass - if dl_dest != dest and dl_dest.exists(): - log(f"Removing {dl_dest}") - dl_dest.unlink() + raise + + if dl_dest != dest and dl_dest.exists(): + log(f"Removing {dl_dest}") + dl_dest.unlink() def api(root_url, service, version, path): diff --git a/test/test_scripts_fetch_content.py b/test/test_scripts_fetch_content.py index 61cfb5b3a..c9f9cb863 100644 --- a/test/test_scripts_fetch_content.py +++ b/test/test_scripts_fetch_content.py @@ -16,6 +16,7 @@ import zipfile from importlib.machinery import SourceFileLoader from importlib.util import module_from_spec, spec_from_loader +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -447,6 +448,196 @@ def test_fetch_urls_merges_staged_extractions( assert (dest / "tests" / "shared.txt").read_text() == "second" +def test_open_for_append_preserves_existing_content(tmp_path, fetch_content_mod): + path = tmp_path / "forward" + path.write_bytes(b"existing\n") + + fd = fetch_content_mod.open_for_append(str(path)) + try: + os.write(fd, b"appended\n") + finally: + os.close(fd) + + assert path.read_bytes() == b"existing\nappended\n" + + +def test_open_for_append_windows(fetch_content_mod, monkeypatch, mocker): + winapi = SimpleNamespace(CreateFile=mocker.Mock(return_value=123), OPEN_EXISTING=3) + msvcrt = SimpleNamespace(open_osfhandle=mocker.Mock(return_value=456)) + monkeypatch.setitem(sys.modules, "_winapi", winapi) + monkeypatch.setitem(sys.modules, "msvcrt", msvcrt) + monkeypatch.setattr(fetch_content_mod, "sys", SimpleNamespace(platform="win32")) + + assert fetch_content_mod.open_for_append("forward") == 456 + + # Append-only access and shared read/write/delete are required for concurrent + # writers; the forward file must already exist. + winapi.CreateFile.assert_called_once_with("forward", 0x4, 0x7, 0, 3, 0, 0) + msvcrt.open_osfhandle.assert_called_once_with(123, 0) + + +@pytest.mark.parametrize("forward", [None, ""]) +def test_record_marker_disabled(fetch_content_mod, monkeypatch, mocker, forward): + monkeypatch.delenv("MOZ_RESOURCE_MONITOR_FORWARD", raising=False) + if forward is not None: + monkeypatch.setenv("MOZ_RESOURCE_MONITOR_FORWARD", forward) + open_for_append = mocker.patch.object(fetch_content_mod, "open_for_append") + + fetch_content_mod.record_marker("fetch-content", 1, 2, "url") + + open_for_append.assert_not_called() + + +@pytest.mark.parametrize("failure", ["open", "write"]) +def test_record_marker_ignores_io_errors( + fetch_content_mod, monkeypatch, mocker, failure +): + monkeypatch.setenv("MOZ_RESOURCE_MONITOR_FORWARD", "forward") + open_for_append = mocker.patch.object( + fetch_content_mod, "open_for_append", return_value=123 + ) + write = mocker.patch.object(fetch_content_mod.os, "write") + close = mocker.patch.object(fetch_content_mod.os, "close") + if failure == "open": + open_for_append.side_effect = OSError("cannot open") + else: + write.side_effect = OSError("cannot write") + + fetch_content_mod.record_marker("fetch-content", 1, 2, "url") + + open_for_append.assert_called_once_with("forward") + if failure == "open": + write.assert_not_called() + close.assert_not_called() + else: + write.assert_called_once_with( + 123, + b'["marker", "fetch-content", 1, 2, {"type": "Text", "text": "url"}, 2]\n', + ) + close.assert_called_once_with(123) + + +def test_fetch_and_extract_records_marker_on_failure(fetch_content_mod, mocker): + error = fetch_content_mod.IntegrityError("invalid checksum") + fetch = mocker.patch.object( + fetch_content_mod, "fetch_and_extract", side_effect=error + ) + record_marker = mocker.patch.object(fetch_content_mod, "record_marker") + mocker.patch.object(fetch_content_mod.time, "monotonic", side_effect=[10, 20]) + + with pytest.raises(fetch_content_mod.IntegrityError) as exc: + fetch_content_mod.fetch_and_extract_with_marker("url", "dest", True, "sha", 42) + + assert exc.value is error + fetch.assert_called_once_with("url", "dest", True, "sha", 42) + record_marker.assert_called_once_with("fetch-content", 10, 20, "url") + + +@pytest.mark.parametrize("command", ["static-url", "git-checkout-archive"]) +@pytest.mark.parametrize("fail", [False, True], ids=["success", "failure"]) +def test_fetch_commands_record_markers( + tmp_path, fetch_content_mod, monkeypatch, mocker, command, fail +): + forward = tmp_path / "forward" + forward.touch() + monkeypatch.setenv("MOZ_RESOURCE_MONITOR_FORWARD", str(forward)) + dest = tmp_path / "fetches" / "archive.tar.zst" + url = ( + "https://example.com/repo" + if command == "git-checkout-archive" + else "https://example.com/archive.tar.gz" + ) + error = RuntimeError("fetch failed") + + def finish(*args, **kwargs): + # The marker must cover the final archive creation, not just the download. + assert forward.read_bytes() == b"" + if fail: + raise error + dest.write_bytes(b"archive") + + if command == "git-checkout-archive": + args = SimpleNamespace( + dest=str(dest), + repo=url, + commit="revision", + path_prefix="repo", + ssh_key_secret=None, + include_dot_git=False, + ) + + def checkout(path, *args, **kwargs): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"partial archive") + finish() + + mocker.patch.object( + fetch_content_mod, "git_checkout_archive", side_effect=checkout + ) + run = fetch_content_mod.command_git_checkout_archive + else: + args = SimpleNamespace( + dest=str(dest), + url=url, + sha256="sha", + size=42, + headers=[], + gpg_sig_url=None, + gpg_key_env=None, + force_archive=False, + strip_components=0, + add_prefix="", + ) + + def download(url, path, **kwargs): + path.write_bytes(b"download") + + mocker.patch.object(fetch_content_mod, "download_to_path", side_effect=download) + mocker.patch.object(fetch_content_mod, "repack_archive", side_effect=finish) + run = fetch_content_mod.command_static_url + + mocker.patch.object(fetch_content_mod.time, "monotonic", side_effect=[10, 20]) + if fail: + with pytest.raises(RuntimeError) as exc: + run(args) + assert exc.value is error + assert not dest.exists() + else: + run(args) + assert dest.read_bytes() == b"archive" + + assert list(dest.parent.iterdir()) == ([] if fail else [dest]) + assert [json.loads(line) for line in forward.read_text().splitlines()] == [ + ["marker", "fetch-content", 10, 20, {"type": "Text", "text": url}, 2] + ] + + +def test_fetch_urls_records_markers( + tmp_path, fetch_content_mod, mock_downloads, monkeypatch +): + forward = tmp_path / "forward" + forward.touch() + monkeypatch.setenv("MOZ_RESOURCE_MONITOR_FORWARD", str(forward)) + dest = tmp_path / "fetches" + dest.mkdir() + mock_downloads["a.txt"] = b"a" + mock_downloads["b.txt"] = b"b" + urls = ["https://example.com/a.txt", "https://example.com/b.txt"] + + fetch_content_mod.fetch_urls([(url, dest, False, None) for url in urls]) + + records = [json.loads(line) for line in forward.read_text().splitlines()] + assert sorted(r[4]["text"] for r in records) == urls + for kind, name, start, end, data, category in records: + assert (kind, name, data["type"], category) == ( + "marker", + "fetch-content", + "Text", + 2, + ) + assert start <= end + + def test_fetch_urls_places_unextracted_files( tmp_path, fetch_content_mod, mock_downloads ):